package com.correx.apps.server.narration

import com.correx.core.approvals.ApprovalOutcome
import com.correx.core.approvals.ApprovalStatus
import com.correx.core.approvals.Tier
import com.correx.core.events.events.ApprovalDecisionResolvedEvent
import com.correx.core.events.events.ApprovalRequestedEvent
import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.EventPayload
import com.correx.core.events.events.ExecutionPlanRejectedEvent
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.OrchestrationPausedEvent
import com.correx.core.events.events.OrchestrationResumedEvent
import com.correx.core.events.events.StageCompletedEvent
import com.correx.core.events.events.StageFailedEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.events.WorkflowCompletedEvent
import com.correx.core.events.events.WorkflowStartedEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ApprovalDecisionId
import com.correx.core.events.types.ApprovalRequestId
import com.correx.core.events.types.EventId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.events.types.TransitionId
import com.correx.core.events.types.ValidationReportId
import org.junit.jupiter.api.Assertions.assertTrue
import com.correx.core.talkie.ChatMode
import com.correx.core.talkie.TalkieFacade
import com.correx.core.talkie.model.NarrationTrigger
import com.correx.core.talkie.model.TalkieResponse
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withTimeout
import kotlinx.coroutines.yield
import kotlinx.datetime.Instant
import org.junit.jupiter.api.AfterEach
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Test
import java.util.concurrent.CopyOnWriteArrayList

class NarrationSubscriberTest {

    private val sessionId = SessionId("session-1")
    private val timestamp = Instant.parse("2026-01-01T00:00:00Z")
    private val scope = CoroutineScope(SupervisorJob() + Dispatchers.Default)

    @AfterEach
    fun tearDown() {
        scope.cancel()
    }

    private val liveFlow = MutableSharedFlow<StoredEvent>(extraBufferCapacity = 64)

    private val fakeStore = object : EventStore {
        override suspend fun append(event: NewEvent): StoredEvent = error("unused")
        override suspend fun appendAll(events: List<NewEvent>): List<StoredEvent> = error("unused")
        override fun read(sessionId: SessionId): List<StoredEvent> = emptyList()
        override fun readFrom(sessionId: SessionId, fromSequence: Long): List<StoredEvent> = emptyList()
        override fun lastSequence(sessionId: SessionId): Long = 0L
        override fun subscribe(sessionId: SessionId): Flow<StoredEvent> = liveFlow
        override fun allEvents(): Sequence<StoredEvent> = emptySequence()
        override fun allSessionIds(): Set<SessionId> = emptySet()
        override fun subscribeAll(): Flow<StoredEvent> = liveFlow
        override suspend fun lastGlobalSequence(): Long = 0L
    }

    private class RecordingTalkieFacade : TalkieFacade {
        val calls = CopyOnWriteArrayList<Pair<SessionId, NarrationTrigger>>()

        override suspend fun onUserInput(sessionId: SessionId, input: String, mode: ChatMode): TalkieResponse =
            error("unused in narration tests")

        override suspend fun narrate(sessionId: SessionId, trigger: NarrationTrigger) {
            calls.add(sessionId to trigger)
        }

        override suspend fun nameSession(sessionId: SessionId, intent: String) = Unit
    }

    private fun storedEvent(payload: EventPayload, seq: Long): StoredEvent = StoredEvent(
        metadata = EventMetadata(
            eventId = EventId("evt-$seq"),
            sessionId = sessionId,
            timestamp = timestamp,
            schemaVersion = 1,
            causationId = null,
            correlationId = null,
        ),
        sequence = seq,
        sessionSequence = seq,
        payload = payload,
    )

    @Test
    fun `triggers are narrated in order for matching events`(): Unit = runBlocking {
        val facade = RecordingTalkieFacade()
        NarrationSubscriber(fakeStore, facade, scope).start()

        // Wait for subscription to be attached
        while (liveFlow.subscriptionCount.value == 0) yield()

        liveFlow.emit(storedEvent(WorkflowStartedEvent(sessionId, "wf", StageId("s0")), seq = 1L))
        liveFlow.emit(storedEvent(StageCompletedEvent(sessionId, StageId("a"), TransitionId("t1")), seq = 2L))
        liveFlow.emit(
            storedEvent(StageFailedEvent(sessionId, StageId("b"), TransitionId("t2"), reason = "x"), seq = 3L),
        )
        liveFlow.emit(storedEvent(WorkflowCompletedEvent(sessionId, StageId("c"), totalStages = 2), seq = 4L))

        withTimeout(2_000L) {
            while (facade.calls.size < 3) yield()
        }

        assertEquals(3, facade.calls.size)
        assertEquals("stage_completed", facade.calls[0].second.kind)
        assertEquals("stage_failed", facade.calls[1].second.kind)
        assertEquals("workflow_completed", facade.calls[2].second.kind)
    }

    @Test
    fun `pause narration is grounded with the pending approval tool and preview`(): Unit = runBlocking {
        val facade = RecordingTalkieFacade()
        NarrationSubscriber(fakeStore, facade, scope).start()

        while (liveFlow.subscriptionCount.value == 0) yield()

        liveFlow.emit(storedEvent(WorkflowStartedEvent(sessionId, "wf", StageId("s0")), seq = 1L))
        liveFlow.emit(
            storedEvent(
                ApprovalRequestedEvent(
                    requestId = ApprovalRequestId("req-1"),
                    tier = Tier.T2,
                    validationReportId = ValidationReportId("vr-1"),
                    riskSummaryId = null,
                    sessionId = sessionId,
                    stageId = StageId("write_script"),
                    projectId = null,
                    toolName = "file_write",
                    preview = "--- a/healthcheck.sh\n+++ b/healthcheck.sh",
                ),
                seq = 2L,
            ),
        )
        liveFlow.emit(
            storedEvent(OrchestrationPausedEvent(sessionId, StageId("write_script"), "APPROVAL_PENDING"), seq = 3L),
        )

        withTimeout(2_000L) { while (facade.calls.isEmpty()) yield() }

        val instruction = facade.calls.first().second.instruction
        assertEquals("paused", facade.calls.first().second.kind)
        assertTrue(instruction.contains("file_write"), "pause narration must name the tool: $instruction")
        assertTrue(instruction.contains("healthcheck.sh"), "pause narration must include the preview: $instruction")
    }

    @Test
    fun `maxPerRun=1 skips second trigger after first narration`(): Unit = runBlocking {
        val facade = RecordingTalkieFacade()
        NarrationSubscriber(fakeStore, facade, scope, maxPerRun = 1).start()

        while (liveFlow.subscriptionCount.value == 0) yield()

        liveFlow.emit(storedEvent(WorkflowStartedEvent(sessionId, "wf", StageId("s0")), seq = 1L))
        liveFlow.emit(storedEvent(StageCompletedEvent(sessionId, StageId("a"), TransitionId("t1")), seq = 2L))
        liveFlow.emit(storedEvent(StageCompletedEvent(sessionId, StageId("b"), TransitionId("t2")), seq = 3L))

        withTimeout(2_000L) {
            while (facade.calls.size < 1) yield()
        }
        // Give a moment for a second call to arrive (it shouldn't)
        kotlinx.coroutines.delay(100L)

        assertEquals(1, facade.calls.size)
        assertEquals("stage_completed", facade.calls[0].second.kind)
    }

    @Test
    fun `a throwing narrate does not prevent subsequent narrations`(): Unit = runBlocking {
        var callCount = 0
        val facade = object : TalkieFacade {
            val calls = CopyOnWriteArrayList<String>()

            override suspend fun onUserInput(sessionId: SessionId, input: String, mode: ChatMode): TalkieResponse =
                error("unused")

            override suspend fun narrate(sessionId: SessionId, trigger: NarrationTrigger) {
                callCount++
                if (callCount == 1) throw RuntimeException("simulated failure")
                calls.add(trigger.kind)
            }

            override suspend fun nameSession(sessionId: SessionId, intent: String) = Unit
        }
        NarrationSubscriber(fakeStore, facade, scope).start()

        while (liveFlow.subscriptionCount.value == 0) yield()

        liveFlow.emit(storedEvent(WorkflowStartedEvent(sessionId, "wf", StageId("s0")), seq = 1L))
        liveFlow.emit(storedEvent(StageCompletedEvent(sessionId, StageId("a"), TransitionId("t1")), seq = 2L))
        liveFlow.emit(storedEvent(StageCompletedEvent(sessionId, StageId("b"), TransitionId("t2")), seq = 3L))

        withTimeout(2_000L) {
            while (facade.calls.size < 1) yield()
        }

        assertEquals(1, facade.calls.size)
        assertEquals("stage_completed", facade.calls[0])
    }

    @Test
    fun `execution plan rejection triggers narration with rejection reason`(): Unit = runBlocking {
        val facade = RecordingTalkieFacade()
        NarrationSubscriber(fakeStore, facade, scope).start()

        while (liveFlow.subscriptionCount.value == 0) yield()

        liveFlow.emit(storedEvent(WorkflowStartedEvent(sessionId, "wf", StageId("s0")), seq = 1L))
        liveFlow.emit(
            storedEvent(
                ExecutionPlanRejectedEvent(sessionId, reason = "Missing required artifact", source = "missing_content"),
                seq = 2L,
            ),
        )

        withTimeout(2_000L) {
            while (facade.calls.isEmpty()) yield()
        }

        assertEquals(1, facade.calls.size)
        assertEquals("plan_rejected", facade.calls[0].second.kind)
        assertTrue(
            facade.calls[0].second.instruction.contains("Missing required artifact"),
            "instruction must include the rejection reason",
        )
        assertTrue(
            facade.calls[0].second.instruction.contains("missing_content"),
            "instruction must include the source",
        )
    }

    /**
     * Router facade that blocks the lane worker on the first narration until [release] completes,
     * letting the test enqueue further events (and resolve the pause) while the pause narration is
     * still pending in the channel.
     */
    private class GatingTalkieFacade : TalkieFacade {
        val calls = CopyOnWriteArrayList<NarrationTrigger>()
        val release = CompletableDeferred<Unit>()
        // Signals that the worker has pulled the first (warmup) trigger and is now blocked,
        // so any narrations enqueued afterwards are guaranteed to still be queued.
        val blocked = CompletableDeferred<Unit>()
        private var first = true

        override suspend fun onUserInput(sessionId: SessionId, input: String, mode: ChatMode): TalkieResponse =
            error("unused in narration tests")

        override suspend fun narrate(sessionId: SessionId, trigger: NarrationTrigger) {
            if (first) {
                first = false
                blocked.complete(Unit)
                release.await()
            }
            calls.add(trigger)
        }

        override suspend fun nameSession(sessionId: SessionId, intent: String) = Unit
    }

    private fun approvalRequested(requestId: String): ApprovalRequestedEvent = ApprovalRequestedEvent(
        requestId = ApprovalRequestId(requestId),
        tier = Tier.T2,
        validationReportId = ValidationReportId("vr-1"),
        riskSummaryId = null,
        sessionId = sessionId,
        stageId = StageId("write_script"),
        projectId = null,
        toolName = "file_write",
        preview = "--- a/x.sh\n+++ b/x.sh",
    )

    private fun approvalResolved(requestId: String): ApprovalDecisionResolvedEvent = ApprovalDecisionResolvedEvent(
        decisionId = ApprovalDecisionId("dec-1"),
        requestId = ApprovalRequestId(requestId),
        outcome = ApprovalOutcome.APPROVED,
        status = ApprovalStatus.COMPLETED,
        tier = Tier.T2,
        resolutionTimestamp = timestamp,
        reason = null,
    )

    @Test
    fun `resolved approval drops the still-queued pause narration`(): Unit = runBlocking {
        val facade = GatingTalkieFacade()
        NarrationSubscriber(fakeStore, facade, scope).start()

        while (liveFlow.subscriptionCount.value == 0) yield()

        liveFlow.emit(storedEvent(WorkflowStartedEvent(sessionId, "wf", StageId("s0")), seq = 1L))
        // Warmup narration the worker will block on, keeping later narrations queued.
        liveFlow.emit(storedEvent(StageCompletedEvent(sessionId, StageId("a"), TransitionId("t1")), seq = 2L))
        // Wait until the worker is parked on the warmup before enqueuing the pause, so the pause
        // is provably still queued (not yet delivered) when its approval resolves.
        withTimeout(2_000L) { facade.blocked.await() }
        // Pause is enqueued, then its approval resolves before the worker can drain it.
        liveFlow.emit(storedEvent(approvalRequested("req-1"), seq = 3L))
        liveFlow.emit(storedEvent(approvalResolved("req-1"), seq = 4L))
        // A current, unrelated narration arrives after resolution and must still surface.
        liveFlow.emit(storedEvent(StageCompletedEvent(sessionId, StageId("b"), TransitionId("t2")), seq = 5L))
        // Let the (single, FIFO) collector drain its buffer so the resolve is handled pre-release.
        delay(100L)

        // Release the worker; it drains the warmup, skips the stale pause, narrates the second stage.
        facade.release.complete(Unit)

        withTimeout(2_000L) { while (facade.calls.size < 2) yield() }
        delay(100L) // allow a stray (stale) pause narration to arrive if the drop failed

        assertEquals(listOf("stage_completed", "stage_completed"), facade.calls.map { it.kind })
        assertTrue(facade.calls.none { it.kind == "paused" }, "stale pause narration must be dropped")
    }

    @Test
    fun `unrelated approval resolution leaves the pending pause intact`(): Unit = runBlocking {
        val facade = GatingTalkieFacade()
        NarrationSubscriber(fakeStore, facade, scope).start()

        while (liveFlow.subscriptionCount.value == 0) yield()

        liveFlow.emit(storedEvent(WorkflowStartedEvent(sessionId, "wf", StageId("s0")), seq = 1L))
        liveFlow.emit(storedEvent(StageCompletedEvent(sessionId, StageId("a"), TransitionId("t1")), seq = 2L))
        withTimeout(2_000L) { facade.blocked.await() }
        // Pause for req-1, but a *different* request (req-2) resolves: req-1's pause must survive.
        liveFlow.emit(storedEvent(approvalRequested("req-1"), seq = 3L))
        liveFlow.emit(storedEvent(approvalResolved("req-2"), seq = 4L))
        delay(100L)

        facade.release.complete(Unit)

        withTimeout(2_000L) { while (facade.calls.size < 2) yield() }

        assertEquals(listOf("stage_completed", "paused"), facade.calls.map { it.kind })
    }

    @Test
    fun `resume drops every queued pause narration for the session`(): Unit = runBlocking {
        val facade = GatingTalkieFacade()
        NarrationSubscriber(fakeStore, facade, scope).start()

        while (liveFlow.subscriptionCount.value == 0) yield()

        liveFlow.emit(storedEvent(WorkflowStartedEvent(sessionId, "wf", StageId("s0")), seq = 1L))
        liveFlow.emit(storedEvent(StageCompletedEvent(sessionId, StageId("a"), TransitionId("t1")), seq = 2L))
        withTimeout(2_000L) { facade.blocked.await() }
        // A non-approval pause (session-scoped) plus an approval pause, both superseded by resume.
        liveFlow.emit(
            storedEvent(OrchestrationPausedEvent(sessionId, StageId("a"), "USER_REQUESTED"), seq = 3L),
        )
        liveFlow.emit(storedEvent(approvalRequested("req-1"), seq = 4L))
        liveFlow.emit(storedEvent(OrchestrationResumedEvent(sessionId, StageId("a")), seq = 5L))
        liveFlow.emit(storedEvent(StageCompletedEvent(sessionId, StageId("b"), TransitionId("t2")), seq = 6L))
        delay(100L)

        facade.release.complete(Unit)

        withTimeout(2_000L) { while (facade.calls.size < 2) yield() }
        delay(100L)

        assertEquals(listOf("stage_completed", "stage_completed"), facade.calls.map { it.kind })
        assertTrue(facade.calls.none { it.kind == "paused" }, "resume must drop all queued pauses")
    }
}
