package com.correx.apps.server.narration

import com.correx.core.approvals.Tier
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.NewEvent
import com.correx.core.events.events.OrchestrationPausedEvent
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.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.router.ChatMode
import com.correx.core.router.RouterFacade
import com.correx.core.router.model.NarrationTrigger
import com.correx.core.router.model.RouterResponse
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
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 RecordingRouterFacade : RouterFacade {
        val calls = CopyOnWriteArrayList<Pair<SessionId, NarrationTrigger>>()

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

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

    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 = RecordingRouterFacade()
        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 = RecordingRouterFacade()
        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 = RecordingRouterFacade()
        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 : RouterFacade {
            val calls = CopyOnWriteArrayList<String>()

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

            override suspend fun narrate(sessionId: SessionId, trigger: NarrationTrigger) {
                callCount++
                if (callCount == 1) throw RuntimeException("simulated failure")
                calls.add(trigger.kind)
            }
        }
        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])
    }
}
