package com.correx.apps.server.logging

import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Test
import org.slf4j.MDC
import kotlin.test.assertEquals
import kotlin.test.assertNull

class CorrelationTest {

    @Test
    fun `withSessionContext sets and restores sessionId in MDC`(): Unit = runBlocking {
        assertNull(MDC.get("sessionId"), "MDC should be empty initially")

        val sessionId = SessionId("test-session-123")
        withSessionContext(sessionId) {
            assertEquals("test-session-123", MDC.get("sessionId"), "sessionId should be set in MDC")
            assertNull(MDC.get("stageId"), "stageId should not be set when not provided")
        }

        assertNull(MDC.get("sessionId"), "sessionId should be restored after block exits")
    }

    @Test
    fun `withSessionContext with stageId sets and restores both values`(): Unit = runBlocking {
        val sessionId = SessionId("session-abc")
        val stageId = StageId("stage-xyz")

        withSessionContext(sessionId, stageId) {
            assertEquals("session-abc", MDC.get("sessionId"), "sessionId should be set")
            assertEquals("stage-xyz", MDC.get("stageId"), "stageId should be set")
        }

        assertNull(MDC.get("sessionId"), "sessionId should be cleared")
        assertNull(MDC.get("stageId"), "stageId should be cleared")
    }

    @Test
    fun `nested withSessionContext stacks and restores correctly`(): Unit = runBlocking {
        val sessionId1 = SessionId("session-1")
        val stageId1 = StageId("stage-1")
        val sessionId2 = SessionId("session-2")
        val stageId2 = StageId("stage-2")

        withSessionContext(sessionId1, stageId1) {
            assertEquals("session-1", MDC.get("sessionId"))
            assertEquals("stage-1", MDC.get("stageId"))

            withSessionContext(sessionId2, stageId2) {
                assertEquals("session-2", MDC.get("sessionId"), "inner context should override")
                assertEquals("stage-2", MDC.get("stageId"), "inner stage should override")
            }

            // After inner block, outer values are restored
            assertEquals("session-1", MDC.get("sessionId"), "outer sessionId should be restored")
            assertEquals("stage-1", MDC.get("stageId"), "outer stageId should be restored")
        }

        assertNull(MDC.get("sessionId"))
        assertNull(MDC.get("stageId"))
    }

    @Test
    fun `withSessionContext clears MDC when overwriting with null stageId`(): Unit = runBlocking {
        val sessionId = SessionId("session-x")

        // Set up initial state
        withSessionContext(sessionId, StageId("initial-stage")) {
            assertEquals("initial-stage", MDC.get("stageId"))
        }

        // Enter new context with same sessionId but no stageId
        withSessionContext(sessionId) {
            assertEquals("session-x", MDC.get("sessionId"))
            assertNull(MDC.get("stageId"), "stageId should be cleared when not provided")
        }

        assertNull(MDC.get("sessionId"))
        assertNull(MDC.get("stageId"))
    }

    @Test
    fun `withSessionContext propagates across coroutine boundaries with MDCContext`(): Unit = runBlocking {
        val sessionId = SessionId("session-coroutine")

        withSessionContext(sessionId) {
            val captured = MDC.get("sessionId")
            assertEquals("session-coroutine", captured, "MDC should be accessible in nested suspension")
        }
    }

    @Test
    fun `withSessionContext removes entries when both sessionId and stageId are absent`(): Unit = runBlocking {
        // Ensure MDC is clean before the test
        MDC.clear()

        withSessionContext(SessionId("new-session")) {
            assertEquals("new-session", MDC.get("sessionId"))
            assertNull(MDC.get("stageId"), "stageId should be cleared")
        }

        assertNull(MDC.get("sessionId"), "sessionId should be cleared after block")
        assertNull(MDC.get("stageId"), "stageId should remain cleared")
    }
}
