package com.correx.apps.server.ws

import com.correx.apps.server.ServerModule
import com.correx.apps.server.configureServer
import com.correx.apps.server.protocol.ServerMessage
import com.correx.apps.server.registry.ProviderRegistry
import com.correx.apps.server.registry.WorkflowRegistry
import com.correx.apps.server.registry.WorkflowSummary
import com.correx.apps.server.undo.SessionUndoService
import com.correx.core.approvals.Tier
import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.events.events.ApprovalRequestedEvent
import com.correx.core.kernel.orchestration.DefaultOrchestrationReducer
import com.correx.core.events.types.ArtifactId
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.events.types.ValidationReportId
import com.correx.core.inference.DefaultInferenceRouter
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.ProviderHealth
import com.correx.core.kernel.orchestration.DefaultSessionOrchestrator
import com.correx.core.kernel.orchestration.OrchestrationConfig
import com.correx.core.kernel.orchestration.OrchestrationProjector
import com.correx.core.kernel.orchestration.OrchestratorEngines
import com.correx.core.kernel.orchestration.OrchestratorRepositories
import com.correx.core.kernel.retry.DefaultRetryCoordinator
import com.correx.core.sessions.DefaultSessionReducer
import com.correx.core.sessions.DefaultSessionRepository
import com.correx.core.sessions.SessionProjector
import com.correx.core.sessions.projections.replay.DefaultEventReplayer
import com.correx.core.transitions.graph.StageConfig
import com.correx.core.transitions.graph.WorkflowGraph
import com.correx.infrastructure.InfrastructureModule
import com.correx.infrastructure.inference.FirstAvailableRoutingStrategy
import com.correx.infrastructure.inference.commons.UnavailableProbe
import com.correx.infrastructure.persistence.InMemoryEventStore
import com.correx.core.kernel.orchestration.OrchestrationRepository
import io.ktor.client.plugins.websocket.WebSockets
import io.ktor.client.plugins.websocket.webSocket
import io.ktor.server.testing.testApplication
import io.ktor.websocket.Frame
import io.ktor.websocket.readText
import kotlinx.coroutines.delay
import kotlinx.coroutines.withTimeout
import kotlinx.serialization.decodeFromString
import kotlinx.serialization.json.Json
import org.junit.jupiter.api.Assertions.assertNotNull
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.io.TempDir
import java.nio.file.Path

/**
 * Covers task #190: a per-session `/sessions/{id}/stream` WS client must observe events that
 * fire *after* connect (e.g. ApprovalRequired for headless CLI --auto-approve), not just the
 * replay-on-connect snapshot. Also covers cleanup on disconnect.
 */
class SessionStreamHandlerTest {

    private val protocolJson = Json { classDiscriminator = "type"; ignoreUnknownKeys = true }
    private fun decode(text: String): ServerMessage = protocolJson.decodeFromString(text)

    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 val noopProviderRegistry: ProviderRegistry = object : ProviderRegistry {
        override fun listAll() = emptyList<InferenceProvider>()
        override suspend fun healthCheckAll() = emptyMap<ProviderId, ProviderHealth>()
    }

    private fun buildWorkflowRegistry(graph: WorkflowGraph): WorkflowRegistry = object : WorkflowRegistry {
        override fun listAll() = listOf(WorkflowSummary(graph.id, description = ""))
        override fun find(workflowId: String) = if (workflowId == graph.id) graph else null
    }

    private fun minimalGraph(): WorkflowGraph = WorkflowGraph(
        id = "session-stream-test-workflow",
        stages = mapOf(StageId("s1") to StageConfig(allowedTools = emptySet())),
        transitions = emptySet(),
        start = StageId("s1"),
    )

    private fun buildModule(tempDir: Path): ServerModule {
        val eventStore = InMemoryEventStore()
        val provider = InfrastructureModule.createLlamaCppProvider(
            modelId = "test-model",
            modelPath = "/dev/null",
            baseUrl = "http://127.0.0.1:1",
        )
        val infraRegistry = InfrastructureModule.createProviderRegistry(listOf(provider))
        val inferenceRouter = DefaultInferenceRouter(infraRegistry, FirstAvailableRoutingStrategy())

        val toolConfig = com.correx.infrastructure.tools.ToolConfig(
            shell = com.correx.infrastructure.tools.ShellConfig(
                enabled = false, allowedExecutables = emptySet(), workingDir = tempDir,
            ),
            fileRead = com.correx.infrastructure.tools.FileReadConfig(
                enabled = false, allowedPaths = setOf(tempDir),
            ),
            fileWrite = com.correx.infrastructure.tools.FileWriteConfig(
                enabled = false, allowedPaths = setOf(tempDir), workingDir = tempDir,
            ),
            fileEdit = com.correx.infrastructure.tools.FileEditConfig(
                enabled = false, allowedPaths = setOf(tempDir), workingDir = tempDir,
            ),
        )
        val toolRegistry = InfrastructureModule.createToolRegistry(toolConfig)
        val eventDispatcher = com.correx.core.events.EventDispatcher(eventStore)
        val toolExecutor = InfrastructureModule.createToolExecutor(
            registry = toolRegistry,
            eventDispatcher = eventDispatcher,
            workDir = tempDir,
            artifactStore = null,
        )

        val engines = OrchestratorEngines(
            transitionResolver = com.correx.core.transitions.resolution.DefaultTransitionResolver {
                    condition, ctx -> condition.evaluate(ctx)
            },
            contextPackBuilder = com.correx.core.context.builder.DefaultContextPackBuilder(
                com.correx.core.context.compression.DefaultContextCompressor(),
            ),
            inferenceRouter = inferenceRouter,
            validationPipeline = com.correx.core.validation.pipeline.ValidationPipeline(validators = emptyList()),
            approvalEngine = com.correx.core.approvals.domain.DefaultApprovalEngine(),
            riskAssessor = com.correx.core.risk.DefaultRiskAssessor(),
            toolRegistry = toolRegistry,
            toolExecutor = toolExecutor,
        )

        val repositories = OrchestratorRepositories(
            eventStore = eventStore,
            inferenceRepository = com.correx.core.inference.InferenceRepository(
                DefaultEventReplayer(eventStore, com.correx.core.inference.InferenceProjector()),
            ),
            orchestrationRepository = OrchestrationRepository(
                DefaultEventReplayer(eventStore, OrchestrationProjector(DefaultOrchestrationReducer())),
            ),
            sessionRepository = DefaultSessionRepository(
                DefaultEventReplayer(eventStore, SessionProjector(DefaultSessionReducer())),
            ),
            artifactRepository = InfrastructureModule.createArtifactRepository(eventStore),
            approvalRepository = InfrastructureModule.createApprovalRepository(eventStore),
        )

        val orchestrator = DefaultSessionOrchestrator(
            repositories = repositories,
            engines = engines,
            retryCoordinator = DefaultRetryCoordinator(eventStore),
            artifactStore = noopArtifactStore,
            tokenizer = provider.tokenizer,
            decisionJournalRepository = InfrastructureModule.createDecisionJournalRepository(eventStore),
        )

        val routerFacade = InfrastructureModule.createTalkieFacade(
            eventStore = eventStore,
            inferenceRouter = inferenceRouter,
            config = com.correx.core.talkie.model.TalkieConfig(),
            tokenizer = provider.tokenizer,
        )

        val sessionUndoService = SessionUndoService(
            eventStore = eventStore,
            artifactStore = noopArtifactStore,
            bootRoots = setOf(tempDir),
        )

        return ServerModule(
            orchestrator = orchestrator,
            eventStore = eventStore,
            artifactStore = noopArtifactStore,
            sessionRepository = repositories.sessionRepository,
            workflowRegistry = buildWorkflowRegistry(minimalGraph()),
            providerRegistry = noopProviderRegistry,
            defaultOrchestrationConfig = OrchestrationConfig(sandboxRoot = tempDir),
            routerFacade = routerFacade,
            orchestrationRepository = repositories.orchestrationRepository,
            approvalRepository = repositories.approvalRepository,
            toolRegistry = toolRegistry,
            sessionUndoService = sessionUndoService,
            resourceProbe = UnavailableProbe,
        )
    }

    @Test
    fun `ApprovalRequired fired after connect reaches the session stream socket`(@TempDir tempDir: Path) {
        val module = buildModule(tempDir)
        module.start()
        val sessionId = SessionId("sid-190")

        testApplication {
            application { configureServer(module) }
            val client = createClient { install(WebSockets) }

            client.webSocket("/sessions/${sessionId.value}/stream") {
                // No lastEventId query param -> no replay, just registerClient() then the read loop.
                delay(100)

                // Fire an ApprovalRequestedEvent through the *same* live fan-out ServerModule.start()
                // wires up (subscribeAll -> approvalCoordinator.onApprovalRequested -> broadcast()),
                // i.e. exactly what a gate firing mid-session would do post-connect.
                module.approvalCoordinator.onApprovalRequested(
                    ApprovalRequestedEvent(
                        requestId = com.correx.core.events.types.ApprovalRequestId("req-190"),
                        tier = Tier.T2,
                        validationReportId = ValidationReportId("vr-190"),
                        riskSummaryId = null,
                        sessionId = sessionId,
                        stageId = null,
                        projectId = null,
                    ),
                )

                val frame = withTimeout(3_000L) {
                    var text: String? = null
                    while (text == null) {
                        val f = incoming.receive()
                        if (f is Frame.Text) text = f.readText()
                    }
                    text
                }
                val decoded = decode(frame)
                assertNotNull(decoded)
                assertTrue(decoded is ServerMessage.ApprovalRequired, "expected ApprovalRequired, got $decoded")
                assertTrue((decoded as ServerMessage.ApprovalRequired).requestId.value == "req-190")
            }
        }
    }

    @Test
    fun `disconnected socket does not block a broadcast to a later-connected sibling`(@TempDir tempDir: Path) {
        // Proves the SessionStreamHandler.handle() finally block actually deregisters the closed
        // socket from ApprovalCoordinator (3559ea67-style cleanup): with two clients on the same
        // session, closing the first and then firing an approval must still deliver cleanly to the
        // second - a stale/dead entry left in sessionClients would surface as a delivery failure
        // log (8d7c827e path) but must never throw or stall the broadcast.
        val module = buildModule(tempDir)
        module.start()
        val sessionId = SessionId("sid-190-disconnect")

        testApplication {
            application { configureServer(module) }
            val client = createClient { install(WebSockets) }

            client.webSocket("/sessions/${sessionId.value}/stream") {
                delay(100)
            }
            // First socket closed on exiting the block; its `finally { unregisterClient(...) }` runs.
            delay(200)

            client.webSocket("/sessions/${sessionId.value}/stream") {
                delay(100)

                module.approvalCoordinator.onApprovalRequested(
                    ApprovalRequestedEvent(
                        requestId = com.correx.core.events.types.ApprovalRequestId("req-190-b"),
                        tier = Tier.T2,
                        validationReportId = ValidationReportId("vr-190-b"),
                        riskSummaryId = null,
                        sessionId = sessionId,
                        stageId = null,
                        projectId = null,
                    ),
                )

                val frame = withTimeout(3_000L) {
                    var text: String? = null
                    while (text == null) {
                        val f = incoming.receive()
                        if (f is Frame.Text) text = f.readText()
                    }
                    text
                }
                val decoded = decode(frame)
                assertTrue(decoded is ServerMessage.ApprovalRequired)
                assertTrue((decoded as ServerMessage.ApprovalRequired).requestId.value == "req-190-b")
            }
        }
    }
}
