package com.correx.apps.server.ws

import com.correx.apps.server.bridge.DomainEventMapper
import com.correx.apps.server.bridge.SessionEventBridge
import com.correx.apps.server.protocol.ServerMessage
import com.correx.apps.server.registry.WorkflowRegistry
import com.correx.core.artifactstore.ArtifactStore
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.StoredEvent
import com.correx.core.events.events.WorkflowStartedEvent
import com.correx.core.events.orchestration.OrchestrationState
import com.correx.core.events.orchestration.OrchestrationStatus
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ArtifactId
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.kernel.orchestration.OrchestrationRepository
import com.correx.core.tools.contract.Tool
import com.correx.core.tools.registry.ToolRegistry
import com.correx.core.sessions.projections.replay.EventReplayer
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.cancelAndJoin
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.launch
import kotlinx.coroutines.test.runTest
import kotlinx.datetime.Instant
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Assertions.assertInstanceOf
import org.junit.jupiter.api.Test

class GlobalStreamHandlerTest {

    private val sessionId = SessionId("session-1")
    private val stageId = StageId("stage-1")
    private val workflowId = "workflow-test"
    private val timestamp = Instant.parse("2026-01-01T00:00:00Z")

    private val noopArtifactStore: ArtifactStore = object : ArtifactStore {
        override suspend fun put(bytes: ByteArray): ArtifactId = ArtifactId("noop")
        override suspend fun get(id: ArtifactId): ByteArray? = null
        override suspend fun flushBefore(commit: suspend () -> Unit) = commit()
    }

    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,
        )

    private fun fakeEventStore(
        liveFlow: MutableSharedFlow<StoredEvent> = MutableSharedFlow(),
        sessions: Set<SessionId> = emptySet(),
        lastGlobal: Long = 0L,
    ): EventStore = 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 = lastGlobal
        override fun subscribe(sessionId: SessionId): Flow<StoredEvent> = liveFlow
        override fun allEvents(): Sequence<StoredEvent> = emptySequence()
        override fun allSessionIds(): Set<SessionId> = sessions
        override fun subscribeAll(): Flow<StoredEvent> = liveFlow
        override suspend fun lastGlobalSequence(): Long = lastGlobal
    }

    private fun activeOrchestrationRepository(): OrchestrationRepository = OrchestrationRepository(
        object : EventReplayer<OrchestrationState> {
            override fun rebuild(sessionId: SessionId): OrchestrationState = OrchestrationState(
                workflowId = workflowId,
                status = OrchestrationStatus.RUNNING,
                currentStageId = stageId,
                pendingApproval = false,
            )
        },
    )

    private fun idleOrchestrationRepository(): OrchestrationRepository = OrchestrationRepository(
        object : EventReplayer<OrchestrationState> {
            override fun rebuild(sessionId: SessionId): OrchestrationState = OrchestrationState(
                workflowId = workflowId,
                status = OrchestrationStatus.IDLE,
                currentStageId = stageId,
                pendingApproval = false,
            )
        },
    )

    private val noopWorkflowRegistry = object : WorkflowRegistry {
        override fun listAll(): List<com.correx.apps.server.registry.WorkflowSummary> = emptyList()
        override fun find(workflowId: String): com.correx.core.transitions.graph.WorkflowGraph? = null
    }

    private val noopToolRegistry = object : ToolRegistry {
        override fun resolve(name: String): Tool? = null
        override fun all(): List<Tool> = emptyList()
    }

    /**
     * Creates a bridge and a unified send channel so that snapshot messages and live messages
     * both flow to the same [received] channel.
     */
    private fun bridgeAndChannel(
        store: EventStore,
        orchRepo: OrchestrationRepository,
    ): Pair<SessionEventBridge, Channel<ServerMessage>> {
        val received = Channel<ServerMessage>(Channel.UNLIMITED)
        val bridge = SessionEventBridge(
            store,
            noopArtifactStore,
            orchRepo,
            noopWorkflowRegistry,
            noopToolRegistry,
        ) { received.send(it) }
        return bridge to received
    }

    @Test
    fun `happy path - snapshot then 3 live events delivered in order`() = runTest {
        val liveFlow = MutableSharedFlow<StoredEvent>(extraBufferCapacity = 16)
        val store = fakeEventStore(liveFlow, setOf(sessionId), lastGlobal = 5L)
        val (bridge, received) = bridgeAndChannel(store, activeOrchestrationRepository())
        val mapper = DomainEventMapper(noopArtifactStore)

        val job = launch {
            // Live messages also go to the same received channel
            streamGlobal(store, bridge, mapper) { received.send(it) }
        }

        delay(50)

        // Emit 3 live events with sequences 6, 7, 8
        repeat(3) { i ->
            liveFlow.emit(storedEvent(WorkflowStartedEvent(sessionId, workflowId, stageId), seq = (6 + i).toLong()))
        }

        delay(100)
        job.cancelAndJoin()
        received.close()

        val messages = mutableListOf<ServerMessage>()
        for (msg in received) messages.add(msg)

        // First: SessionSnapshot (from active session)
        assertInstanceOf(ServerMessage.SessionSnapshot::class.java, messages[0])
        // Second: SnapshotComplete
        assertEquals(ServerMessage.SnapshotComplete, messages[1])
        // Then 3 WorkflowStarted messages (seq 6, 7, 8)
        val live = messages.drop(2)
        assertEquals(3, live.size)
        live.forEach { assertInstanceOf(ServerMessage.SessionAnnounced::class.java, it) }
    }

    @Test
    fun `race - event emitted during replaySnapshot is not lost`() = runTest {
        repeat(100) {
            val liveFlow = MutableSharedFlow<StoredEvent>(extraBufferCapacity = 64)
            val store = fakeEventStore(liveFlow, emptySet(), lastGlobal = 5L)
            val received = Channel<ServerMessage>(Channel.UNLIMITED)
            val bridge = SessionEventBridge(
                store,
                noopArtifactStore,
                idleOrchestrationRepository(),
                noopWorkflowRegistry,
                noopToolRegistry,
            ) {
                received.send(it)
            }
            val mapper = DomainEventMapper(noopArtifactStore)
            val raceEvent = storedEvent(WorkflowStartedEvent(sessionId, workflowId, stageId), seq = 6L)

            // Emit the race event concurrently with streamGlobal starting.
            // streamGlobal subscribes BEFORE calling replaySnapshot, so this event will be buffered
            // and delivered after SnapshotComplete even if it arrives during the snapshot phase.
            val emitter = launch {
                delay(1)
                liveFlow.emit(raceEvent)
            }

            val job = launch {
                streamGlobal(store, bridge, mapper) { received.send(it) }
            }

            delay(100)
            emitter.cancelAndJoin()
            job.cancelAndJoin()
            received.close()

            val messages = mutableListOf<ServerMessage>()
            for (msg in received) messages.add(msg)

            assert(messages.contains(ServerMessage.SnapshotComplete)) {
                "SnapshotComplete missing in iteration $it"
            }
            val completeIdx = messages.indexOf(ServerMessage.SnapshotComplete)
            val afterComplete = messages.drop(completeIdx + 1)
            assert(afterComplete.any { msg -> msg is ServerMessage.SessionAnnounced }) {
                "Race event (seq=6) was lost in iteration $it! Messages: $messages"
            }
        }
    }

    @Test
    fun `empty - zero sessions - receives SnapshotComplete then live events`() = runTest {
        val liveFlow = MutableSharedFlow<StoredEvent>(extraBufferCapacity = 16)
        val store = fakeEventStore(liveFlow, emptySet(), lastGlobal = 0L)
        val received = Channel<ServerMessage>(Channel.UNLIMITED)
        val bridge = SessionEventBridge(
            store,
            noopArtifactStore,
            idleOrchestrationRepository(),
            noopWorkflowRegistry,
            noopToolRegistry,
        ) {
            received.send(it)
        }
        val mapper = DomainEventMapper(noopArtifactStore)

        val job = launch {
            streamGlobal(store, bridge, mapper) { received.send(it) }
        }

        delay(50)

        // Emit a live event after snapshot
        liveFlow.emit(storedEvent(WorkflowStartedEvent(sessionId, workflowId, stageId), seq = 1L))

        delay(50)
        job.cancelAndJoin()
        received.close()

        val messages = mutableListOf<ServerMessage>()
        for (msg in received) messages.add(msg)

        assertEquals(ServerMessage.SnapshotComplete, messages[0])
        assertEquals(2, messages.size)
        assertInstanceOf(ServerMessage.SessionAnnounced::class.java, messages[1])
    }

    @Test
    fun `barrier - subscription is registered before lastGlobalSequence is read`() = runTest {
        val liveFlow = MutableSharedFlow<StoredEvent>(extraBufferCapacity = 16)
        val observedAtSnapshotRead = CompletableDeferred<Int>()

        val store = 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 {
                // Capture the subscription count at the moment replaySnapshot reads the offset.
                // The barrier in streamGlobal must guarantee count > 0 here.
                observedAtSnapshotRead.complete(liveFlow.subscriptionCount.value)
                return 0L
            }
        }
        val bridge = SessionEventBridge(
            store,
            noopArtifactStore,
            idleOrchestrationRepository(),
            noopWorkflowRegistry,
            noopToolRegistry,
        ) { }
        val mapper = DomainEventMapper(noopArtifactStore)

        val job = launch {
            streamGlobal(store, bridge, mapper) { }
        }

        val count = observedAtSnapshotRead.await()
        job.cancelAndJoin()

        assertEquals(1, count, "subscribeAll() must have an active subscriber before lastGlobalSequence() is read")
    }

    @Test
    fun `shutdown - cancel outer scope ends subscription with no leaks`() = runTest {
        val liveFlow = MutableSharedFlow<StoredEvent>()
        val store = fakeEventStore(liveFlow, emptySet(), lastGlobal = 0L)
        val bridge = SessionEventBridge(
            store,
            noopArtifactStore,
            idleOrchestrationRepository(),
            noopWorkflowRegistry,
            noopToolRegistry,
        ) { }
        val mapper = DomainEventMapper(noopArtifactStore)

        val job = launch {
            streamGlobal(store, bridge, mapper) { }
        }

        delay(50)
        job.cancelAndJoin()

        assertFalse(job.isActive, "Handler job must not be active after cancel")
    }
}
