import com.correx.core.context.model.ContextLayer
import com.correx.core.context.model.ContextPack
import com.correx.core.context.model.TokenBudget
import com.correx.core.events.events.ChatTurnEvent
import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.RouterNarrationEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ContextPackId
import com.correx.core.events.types.EventId
import com.correx.core.events.types.ProviderId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.inference.CapabilityScore
import com.correx.core.inference.Embedder
import com.correx.core.inference.FinishReason
import com.correx.core.inference.PromptRenderer
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.InferenceResponse
import com.correx.core.inference.InferenceRouter
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.NoopEmbedder
import com.correx.core.inference.ProviderHealth
import com.correx.core.inference.ResponseFormat
import com.correx.core.inference.TokenUsage
import com.correx.core.router.DefaultRouterContextBuilder
import com.correx.core.router.DefaultRouterFacade
import com.correx.core.router.RouterContextBuilder
import com.correx.core.router.RouterRepository
import com.correx.core.router.l3.InMemoryL3MemoryStore
import com.correx.core.router.l3.L3MemoryEntry
import com.correx.core.router.l3.L3MemoryStore
import com.correx.core.router.l3.L3Query
import com.correx.core.router.model.NarrationTrigger
import com.correx.core.router.model.RouterConfig
import com.correx.core.router.model.RouterState
import com.correx.core.router.model.WorkflowStatus
import com.correx.core.router.model.WorkflowSummary
import com.correx.testing.fixtures.inference.MockTokenizer
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertNotNull
import org.junit.jupiter.api.Assertions.assertNull
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test

class RouterNarrationTest {

    // --------------------------------------------------------------------------
    // narrate() — event emission
    // --------------------------------------------------------------------------

    @Test
    fun `narrate emits RouterNarrationEvent with inference content and metrics`(): Unit = runBlocking {
        val knownLatencyMs = 55L
        val knownUsage = TokenUsage(promptTokens = 8, completionTokens = 3)
        val eventStore = MapBackedEventStore()
        val facade = buildFacade(
            eventStore = eventStore,
            responseText = "narration line",
            latencyMs = knownLatencyMs,
            tokensUsed = knownUsage,
        )
        val trigger = NarrationTrigger(kind = "stage_completed", instruction = "Summarise what just happened.")

        facade.narrate(SessionId("sess-1"), trigger)

        val narrationEvents = eventStore.appendedEvents.map { it.payload }.filterIsInstance<RouterNarrationEvent>()
        assertEquals(1, narrationEvents.size)
        val event = narrationEvents[0]
        assertEquals("narration line", event.content)
        assertEquals(knownLatencyMs, event.latencyMs)
        assertNotNull(event.tokensUsed)
        assertEquals(11, event.tokensUsed!!.totalTokens)
        assertEquals("stage_completed", event.trigger)
    }

    @Test
    fun `narrate does not emit ChatTurnEvent`(): Unit = runBlocking {
        val eventStore = MapBackedEventStore()
        val facade = buildFacade(eventStore = eventStore, responseText = "narration text")
        val trigger = NarrationTrigger(kind = "tool_executed", instruction = "Describe the tool result.")

        facade.narrate(SessionId("sess-2"), trigger)

        val chatEvents = eventStore.appendedEvents.map { it.payload }.filterIsInstance<ChatTurnEvent>()
        assertTrue(chatEvents.isEmpty(), "narrate must not emit any ChatTurnEvent")
    }

    @Test
    fun `narrate does not write to l3MemoryStore`(): Unit = runBlocking {
        val trackingL3 = TrackingL3MemoryStore()
        val facade = buildFacade(
            l3MemoryStore = trackingL3,
            responseText = "narration output",
        )
        val trigger = NarrationTrigger(kind = "stage_completed", instruction = "Tell the user.")

        facade.narrate(SessionId("sess-3"), trigger)

        assertTrue(trackingL3.storedEntries.isEmpty(), "narrate must not call l3MemoryStore.store")
    }

    @Test
    fun `narrate skips event emission when inference returns blank`(): Unit = runBlocking {
        val eventStore = MapBackedEventStore()
        val facade = buildFacade(eventStore = eventStore, responseText = "   ")
        val trigger = NarrationTrigger(kind = "stage_completed", instruction = "Say something.")

        facade.narrate(SessionId("sess-4"), trigger)

        val narrationEvents = eventStore.appendedEvents.map { it.payload }.filterIsInstance<RouterNarrationEvent>()
        assertTrue(narrationEvents.isEmpty(), "blank inference must not emit RouterNarrationEvent")
    }

    @Test
    fun `narrate skips event emission when inference returns empty string`(): Unit = runBlocking {
        val eventStore = MapBackedEventStore()
        val facade = buildFacade(eventStore = eventStore, responseText = "")
        val trigger = NarrationTrigger(kind = "tool_executed", instruction = "Narrate.")

        facade.narrate(SessionId("sess-5"), trigger)

        val narrationEvents = eventStore.appendedEvents.map { it.payload }.filterIsInstance<RouterNarrationEvent>()
        assertTrue(narrationEvents.isEmpty(), "empty inference must not emit RouterNarrationEvent")
    }

    @Test
    fun `narrate RouterNarrationEvent carries sessionId and stageId from state`(): Unit = runBlocking {
        val stageId = StageId("my-stage")
        val eventStore = MapBackedEventStore()
        val facade = DefaultRouterFacade(
            routerRepository = object : RouterRepository {
                override suspend fun getRouterState(sessionId: SessionId): RouterState =
                    RouterState(
                        sessionId = sessionId,
                        workflowStatus = WorkflowStatus.RUNNING,
                        currentStageId = stageId,
                    )
            },
            routerContextBuilder = object : RouterContextBuilder {
                override suspend fun build(state: RouterState, budget: TokenBudget, availableWorkflows: List<WorkflowSummary>, projectProfileText: String?): ContextPack = emptyContextPack()
                override suspend fun buildNarrationContext(state: RouterState, trigger: NarrationTrigger, budget: TokenBudget): ContextPack = emptyContextPack()
            },
            inferenceRouter = mockInferenceRouter("stage narration", latencyMs = 10L, tokensUsed = TokenUsage(1, 1)),
            eventStore = eventStore,
            config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
            embedder = NoopEmbedder(dimension = 8),
            l3MemoryStore = InMemoryL3MemoryStore(),
        )
        val trigger = NarrationTrigger(kind = "stage_completed", instruction = "Describe the stage.")

        facade.narrate(SessionId("session-x"), trigger)

        val narrationEvents = eventStore.appendedEvents.map { it.payload }.filterIsInstance<RouterNarrationEvent>()
        assertEquals(1, narrationEvents.size)
        assertEquals(SessionId("session-x"), narrationEvents[0].sessionId)
        assertEquals("my-stage", narrationEvents[0].stageId)
    }

    // --------------------------------------------------------------------------
    // buildNarrationContext — content structure
    // --------------------------------------------------------------------------

    @Test
    fun `buildNarrationContext includes system entry in L0`(): Unit = runBlocking {
        val builder = DefaultRouterContextBuilder(RouterConfig())
        val state = RouterState(sessionId = SessionId("s"), workflowStatus = WorkflowStatus.IDLE)
        val trigger = NarrationTrigger(kind = "test", instruction = "Do something.")
        val pack = builder.buildNarrationContext(state, trigger, TokenBudget(limit = 5000))

        val l0 = pack.layers[ContextLayer.L0]
        assertNotNull(l0, "L0 layer must be present")
        assertTrue(l0!!.isNotEmpty(), "L0 must contain at least the system prompt")
    }

    @Test
    fun `buildNarrationContext includes workflow status entry`(): Unit = runBlocking {
        val builder = DefaultRouterContextBuilder(RouterConfig())
        val state = RouterState(
            sessionId = SessionId("s"),
            workflowStatus = WorkflowStatus.RUNNING,
            currentStageId = StageId("my-stage"),
        )
        val trigger = NarrationTrigger(kind = "test", instruction = "Describe.")
        val pack = builder.buildNarrationContext(state, trigger, TokenBudget(limit = 5000))

        val l0 = pack.layers[ContextLayer.L0]
        assertNotNull(l0)
        val allContent = l0!!.joinToString(" ") { it.content }
        assertTrue(allContent.contains("my-stage"), "workflow status must mention the current stage")
    }

    @Test
    fun `buildNarrationContext includes trigger instruction as final entry`(): Unit = runBlocking {
        val builder = DefaultRouterContextBuilder(RouterConfig())
        val state = RouterState(sessionId = SessionId("s"), workflowStatus = WorkflowStatus.IDLE)
        val instruction = "Narrate the completed action now."
        val trigger = NarrationTrigger(kind = "stage_completed", instruction = instruction)
        val pack = builder.buildNarrationContext(state, trigger, TokenBudget(limit = 5000))

        val allEntries = pack.layers.values.flatten()
        assertTrue(
            allEntries.any { it.content.contains(instruction) },
            "trigger instruction must appear in narration context",
        )
    }

    @Test
    fun `rendered narration prompt carries a user message`(): Unit = runBlocking {
        // Regression: the trigger instruction must render as a user turn. When it was filed
        // under L0 it folded into the system message, leaving zero user messages and the chat
        // template rejected the request ("No user query found in messages").
        val builder = DefaultRouterContextBuilder(RouterConfig())
        val state = RouterState(sessionId = SessionId("s"), workflowStatus = WorkflowStatus.RUNNING)
        val instruction = "Tell the user the stage just completed."
        val trigger = NarrationTrigger(kind = "stage_completed", instruction = instruction)
        val pack = builder.buildNarrationContext(state, trigger, TokenBudget(limit = 5000))

        val messages = PromptRenderer.render(pack)
        val userMessages = messages.filter { it.role == "user" }
        assertTrue(userMessages.isNotEmpty(), "narration prompt must contain a user message")
        assertTrue(
            userMessages.any { it.content.contains(instruction) },
            "the trigger instruction must be the user turn",
        )
    }

    @Test
    fun `buildNarrationContext stays within budget`(): Unit = runBlocking {
        val builder = DefaultRouterContextBuilder(RouterConfig())
        val state = RouterState(sessionId = SessionId("s"), workflowStatus = WorkflowStatus.IDLE)
        val trigger = NarrationTrigger(kind = "test", instruction = "Short.")
        val budget = TokenBudget(limit = 5000)
        val pack = builder.buildNarrationContext(state, trigger, budget)

        assertTrue(pack.budgetUsed <= budget.limit, "narration context must not exceed budget")
    }

    @Test
    fun `buildNarrationContext does not include conversationHistory turns`(): Unit = runBlocking {
        val builder = DefaultRouterContextBuilder(RouterConfig())
        val state = RouterState(
            sessionId = SessionId("s"),
            workflowStatus = WorkflowStatus.IDLE,
            conversationHistory = listOf(
                com.correx.core.router.model.RouterTurn(
                    role = com.correx.core.router.model.TurnRole.USER,
                    content = "secret conversation turn",
                    timestamp = kotlinx.datetime.Clock.System.now(),
                ),
            ),
        )
        val trigger = NarrationTrigger(kind = "test", instruction = "Narrate.")
        val pack = builder.buildNarrationContext(state, trigger, TokenBudget(limit = 5000))

        val allContent = pack.layers.values.flatten().joinToString(" ") { it.content }
        assertTrue(
            !allContent.contains("secret conversation turn"),
            "buildNarrationContext must not include conversationHistory turns",
        )
    }

    // --------------------------------------------------------------------------
    // Helpers
    // --------------------------------------------------------------------------

    private fun buildFacade(
        eventStore: EventStore = MapBackedEventStore(),
        responseText: String = "narration",
        latencyMs: Long = 10L,
        tokensUsed: TokenUsage = TokenUsage(5, 5),
        l3MemoryStore: L3MemoryStore = InMemoryL3MemoryStore(),
    ): DefaultRouterFacade = DefaultRouterFacade(
        routerRepository = object : RouterRepository {
            override suspend fun getRouterState(sessionId: SessionId): RouterState =
                RouterState(
                    sessionId = sessionId,
                    workflowStatus = WorkflowStatus.RUNNING,
                    currentStageId = StageId("s1"),
                )
        },
        routerContextBuilder = DefaultRouterContextBuilder(RouterConfig()),
        inferenceRouter = mockInferenceRouter(responseText, latencyMs, tokensUsed),
        eventStore = eventStore,
        config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
        embedder = NoopEmbedder(dimension = 8),
        l3MemoryStore = l3MemoryStore,
    )

    private fun mockInferenceRouter(
        responseText: String,
        latencyMs: Long = 10L,
        tokensUsed: TokenUsage = TokenUsage(5, 5),
    ): InferenceRouter = object : InferenceRouter {
        override suspend fun route(
            stageId: StageId,
            requiredCapabilities: Set<ModelCapability>,
        ): InferenceProvider = object : InferenceProvider {
            override val id = ProviderId("mock")
            override val name = "Mock"
            override val tokenizer = MockTokenizer()
            override suspend fun infer(request: InferenceRequest): InferenceResponse = InferenceResponse(
                requestId = request.requestId,
                text = responseText,
                finishReason = FinishReason.Stop,
                tokensUsed = tokensUsed,
                latencyMs = latencyMs,
            )
            override suspend fun healthCheck(): ProviderHealth = ProviderHealth.Healthy
            override fun capabilities(): Set<CapabilityScore> = setOf(CapabilityScore(ModelCapability.General, 1.0))
        }
    }

    private fun emptyContextPack(): ContextPack = ContextPack(
        id = ContextPackId("empty"),
        sessionId = SessionId("unknown"),
        stageId = StageId("none"),
        layers = emptyMap(),
        budgetUsed = 0,
        budgetLimit = 5000,
    )

    private class MapBackedEventStore : EventStore {
        val appendedEvents = mutableListOf<NewEvent>()
        private val storedEvents: MutableMap<EventId, StoredEvent> = mutableMapOf()
        private var nextSequence = 1L

        override suspend fun append(event: NewEvent): StoredEvent {
            appendedEvents.add(event)
            val stored = StoredEvent(
                metadata = event.metadata,
                sequence = nextSequence,
                sessionSequence = nextSequence++,
                payload = event.payload,
            )
            storedEvents[event.metadata.eventId] = stored
            return stored
        }

        override suspend fun appendAll(events: List<NewEvent>): List<StoredEvent> =
            events.map { append(it) }

        override fun read(sessionId: SessionId): List<StoredEvent> =
            storedEvents.values.filter { it.metadata.sessionId == sessionId }.toList()

        override fun readFrom(sessionId: SessionId, fromSequence: Long): List<StoredEvent> =
            read(sessionId).filter { it.sequence >= fromSequence }

        override fun lastSequence(sessionId: SessionId): Long? =
            read(sessionId).maxOfOrNull { it.sequence }

        override fun subscribe(sessionId: SessionId): Flow<StoredEvent> =
            throw UnsupportedOperationException("subscribe not implemented for mock")

        override fun allEvents(): Sequence<StoredEvent> = storedEvents.values.asSequence()

        override fun allSessionIds(): Set<SessionId> = storedEvents.values.map { it.metadata.sessionId }.toSet()

        override fun subscribeAll(): Flow<StoredEvent> = TODO("Not needed in this test context")

        override suspend fun lastGlobalSequence(): Long = TODO("Not needed in this test context")
    }

    private class TrackingL3MemoryStore : L3MemoryStore {
        val storedEntries = mutableListOf<L3MemoryEntry>()

        override suspend fun store(entry: L3MemoryEntry) {
            storedEntries.add(entry)
        }

        override suspend fun query(query: L3Query): List<com.correx.core.router.l3.L3Hit> = emptyList()

        override suspend fun existsByTurnIdPrefix(prefix: String): Boolean = false

        override suspend fun close() = Unit
    }
}