package com.correx.apps.server.ws

import com.correx.apps.server.ServerModule
import com.correx.apps.server.bridge.DomainEventMapper
import com.correx.apps.server.bridge.SessionEventBridge
import com.correx.apps.server.logging.withSessionContext
import com.correx.apps.server.protocol.ClientMessage
import com.correx.apps.server.protocol.GrantScopeDto
import com.correx.apps.server.protocol.ProtocolSerializer
import com.correx.apps.server.protocol.ProviderHealthDto
import com.correx.apps.server.protocol.ServerMessage
import com.correx.apps.server.protocol.WorkflowDto
import com.correx.apps.server.protocol.StageToolDecl
import com.correx.apps.server.protocol.ToolDecl
import com.correx.apps.server.workspace.WorkspaceResolution
import com.correx.core.approvals.GrantScope
import com.correx.core.approvals.Tier
import com.correx.core.events.events.ApprovalGrantCreatedEvent
import com.correx.core.events.events.ChatSessionStartedEvent
import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.IdeaDiscardedEvent
import com.correx.core.events.events.InitialIntentEvent
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.SessionWorkspaceBoundEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.types.GrantId
import com.correx.core.events.stores.EventStore
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.inference.ProviderHealth
import com.correx.core.kernel.orchestration.WorkspaceContext
import com.correx.core.router.IdeaReader
import com.correx.core.utils.TypeId
import com.correx.infrastructure.inference.commons.ResourceProbe
import com.correx.infrastructure.inference.commons.ResourceSnapshot
import io.ktor.server.websocket.DefaultWebSocketServerSession
import io.ktor.websocket.Frame
import io.ktor.websocket.readText
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.channels.BufferOverflow
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.ClosedReceiveChannelException
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.SharedFlow
import kotlinx.coroutines.flow.onStart
import kotlinx.coroutines.flow.onSubscription
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import kotlinx.datetime.Clock
import org.slf4j.LoggerFactory
import java.util.UUID

private val log = LoggerFactory.getLogger(GlobalStreamHandler::class.java)
private const val BUFFER_CAPACITY = 1024
private const val RESOURCE_PUSH_INTERVAL_MS = 2500L
private const val BYTES_PER_MB = 1024L * 1024L

private fun ResourceSnapshot.toResourceStatus(): ServerMessage.ResourceStatus =
    ServerMessage.ResourceStatus(
        gpuMemoryUsedMb = gpu?.memoryUsedMb,
        gpuMemoryTotalMb = gpu?.memoryTotalMb,
        gpuUtilizationPct = gpu?.utilizationPct,
        processRssMb = processRssBytes?.let { it / BYTES_PER_MB },
        systemRamUsedMb = systemRamUsedBytes?.let { it / BYTES_PER_MB },
        systemRamTotalMb = systemRamTotalBytes?.let { it / BYTES_PER_MB },
    )

class GlobalStreamHandler(private val module: ServerModule) {

    private val queries = StreamQueries(module)

    suspend fun handle(session: DefaultWebSocketServerSession) = coroutineScope {
        log.info("client connected")

        val sendFrame: suspend (ServerMessage) -> Unit = {
            session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(it)))
        }

        // Stash the working directory received in the Hello frame (first frame on connect).
        // Null means no Hello was received; handleStartSession falls back to boot default.
        var pendingWorkingDir: String? = null
        // Guard: once StartSession has been handled the Hello window is closed. A late Hello
        // has no session to bind to and its workingDir is silently stale — ignore it.
        var sessionStarted = false

        val bridge = SessionEventBridge(
            eventStore = module.eventStore,
            artifactStore = module.artifactStore,
            orchestrationRepository = module.orchestrationRepository,
            workflowRegistry = module.workflowRegistry,
            toolRegistry = module.toolRegistry,
            send = sendFrame,
            approvalCoordinator = module.approvalCoordinator,
        )
        val mapper = DomainEventMapper(module.artifactStore)

        module.approvalCoordinator.registerGlobalClient(session)

        sendInitialSnapshot(session)

        // Launch global stream (snapshot + live-forward) concurrently with incoming frame loop
        launch {
            streamGlobal(module.eventStore, bridge, mapper, sendFrame)
        }
        // Live resource gauge (GPU/RAM) pushed periodically. Cancelled with the session scope.
        launch {
            streamResources(module.resourceProbe, sendFrame)
        }

        try {
            for (frame in session.incoming) {
                if (frame is Frame.Text) {
                    val text = frame.readText()
                    runCatching { ProtocolSerializer.decodeClientMessage(text) }
                        .onSuccess { msg ->
                            log.debug("recv: {}", msg::class.simpleName)
                            runCatching {
                                handleClientMessage(
                                    session, msg, sendFrame, pendingWorkingDir, sessionStarted,
                                    onHello = { pendingWorkingDir = it },
                                    onSessionStarted = { sessionStarted = true },
                                )
                            }.onFailure {
                                log.error("handle failed", it)
                            }
                        }
                        .onFailure {
                            log.warn("decode error: {}", it.message)
                            val error = ServerMessage.ProtocolError("Unknown message: ${it.message}")
                            session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(error)))
                        }
                }
            }
        } catch (_: ClosedReceiveChannelException) {
            log.info("client disconnected")
        } finally {
            module.approvalCoordinator.unregisterGlobalClient(session)
        }
    }

    private suspend fun sendInitialSnapshot(session: DefaultWebSocketServerSession) {
        val providerHealth = runCatching { module.providerRegistry.healthCheckAll() }
            .onFailure { log.warn("healthCheckAll failed: {}", it.message) }
            .getOrDefault(emptyMap())
        log.info("snapshot: {} provider(s)", providerHealth.size)
        providerHealth.forEach { (providerId, health) ->
            val status = when (health) {
                is ProviderHealth.Healthy -> "healthy"
                is ProviderHealth.Degraded -> "degraded"
                is ProviderHealth.Unavailable -> "unavailable"
            }
            log.debug("provider={} status={}", providerId.value, status)
            val msg = ServerMessage.ProviderStatusChanged(
                providerId = providerId.value,
                status = ProviderHealthDto(providerId.value, status, null),
                sequence = null,
                sessionSequence = null,
            )
            session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(msg)))
        }

        // Send available workflows so the TUI can show a workflow picker
        val workflows = module.workflowRegistry.listAll().map { summary ->
            WorkflowDto(workflowId = summary.workflowId, description = summary.description)
        }
        log.debug("workflow list: {} workflow(s)", workflows.size)
        session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(
            ServerMessage.WorkflowList(workflows = workflows),
        )))

        // Configured models + current, so the client can offer a swap picker (managed path only).
        module.modelSwapper?.let { swapper ->
            session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(
                ServerMessage.ModelList(
                    models = swapper.availableModelIds(),
                    current = swapper.currentModelId(),
                ),
            )))
        }

        // Initial resource gauge so a freshly-connected client renders VRAM/RAM immediately.
        val snapshot = withContext(Dispatchers.IO) { module.resourceProbe.probe() }
        session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(snapshot.toResourceStatus())))
    }

    /** Pushes a [ServerMessage.ResourceStatus] every [RESOURCE_PUSH_INTERVAL_MS]. */
    private suspend fun streamResources(
        probe: ResourceProbe,
        sendFrame: suspend (ServerMessage) -> Unit,
    ) {
        while (true) {
            delay(RESOURCE_PUSH_INTERVAL_MS)
            val snapshot = withContext(Dispatchers.IO) { probe.probe() }
            sendFrame(snapshot.toResourceStatus())
        }
    }

    // Flat dispatch table: an exhaustive `when` over the sealed ClientMessage. Its cyclomatic
    // complexity scales with the number of message types, which is the false-positive case the
    // metric flags — each branch is a one-line delegation, not real branching logic.
    @Suppress("CyclomaticComplexMethod")
    private suspend fun handleClientMessage(
        session: DefaultWebSocketServerSession,
        msg: ClientMessage,
        sendFrame: suspend (ServerMessage) -> Unit,
        pendingWorkingDir: String?,
        sessionStarted: Boolean,
        onHello: (String) -> Unit,
        onSessionStarted: () -> Unit,
    ) {
        when (msg) {
            is ClientMessage.Ping -> Unit
            is ClientMessage.Hello -> {
                if (sessionStarted) {
                    log.warn("recv Hello after session already started (workingDir={}); ignoring", msg.workingDir)
                } else {
                    log.info("recv Hello: workingDir={}", msg.workingDir)
                    onHello(msg.workingDir)
                }
            }
            is ClientMessage.StartSession -> {
                handleStartSession(session, msg, pendingWorkingDir)
                onSessionStarted()
            }
            is ClientMessage.StartChatSession -> handleStartChatSession(session, msg, sendFrame)
            is ClientMessage.CancelSession -> {
                runCatching { module.orchestrator.cancel(msg.sessionId) }
                    .onFailure { log.error("cancel failed for session={}: {}", msg.sessionId.value, it.message, it) }
            }
            is ClientMessage.ListArtifacts -> sendFrame(queries.listArtifacts(msg.sessionId))
            is ClientMessage.ListIdeas -> sendFrame(queries.listIdeas())
            is ClientMessage.DiscardIdea -> handleDiscardIdea(msg, sendFrame)
            is ClientMessage.GetSessionStats -> sendFrame(queries.sessionStats(msg.sessionId))
            is ClientMessage.GetConfig -> sendFrame(queries.configSnapshot(error = null, restartRequired = emptyList()))
            is ClientMessage.UpdateConfig -> sendFrame(queries.updateConfig(msg.patch))
            is ClientMessage.ResumeSession -> session.send(encodeError("ResumeSession not supported"))
            is ClientMessage.ApprovalResponse -> handleApprovalResponse(msg, sendFrame)
            is ClientMessage.ClarificationResponse -> handleClarificationResponse(msg)
            is ClientMessage.CreateGrant -> handleCreateGrant(msg, sendFrame)
            is ClientMessage.ChatInput -> {
                // The USER and ROUTER turns are emitted as ChatTurnEvents inside onUserInput
                // and reach the client as chat.turn frames via streamGlobal — no direct
                // RouterResponseMessage (the conversation is event-derived).
                withSessionContext(msg.sessionId) {
                    runCatching {
                        module.routerFacade.onUserInput(msg.sessionId, msg.text, msg.mode)
                    }.onFailure {
                        log.error("routerFacade.onUserInput failed: {}", it.message)
                        sendFrame(ServerMessage.ProtocolError("Router error: ${it.message}"))
                    }
                }
            }
            is ClientMessage.SwapModel -> handleSwapModel(msg, sendFrame)
            is ClientMessage.ClearModelPin -> {
                val swapper = module.modelSwapper
                if (swapper == null) {
                    sendFrame(errorResponse("Model management is not enabled"))
                } else {
                    swapper.clearPin()
                    log.info("model pin cleared")
                }
            }
        }
    }

    private suspend fun handleSwapModel(
        msg: ClientMessage.SwapModel,
        sendFrame: suspend (ServerMessage) -> Unit,
    ) {
        val swapper = module.modelSwapper
        if (swapper == null) {
            sendFrame(errorResponse("Model management is not enabled"))
            return
        }
        // swap() makes the model resident; the resulting ModelLoadedEvent surfaces the change to
        // clients via streamGlobal, so we don't push ModelChanged here.
        runCatching { swapper.swap(msg.modelId) }
            .onSuccess { log.info("swapped to model={}", msg.modelId) }
            .onFailure {
                log.error("model swap failed for '{}': {}", msg.modelId, it.message)
                sendFrame(errorResponse("Model swap failed: ${it.message}"))
            }
    }

    private suspend fun handleApprovalResponse(
        msg: ClientMessage.ApprovalResponse,
        sendFrame: suspend (ServerMessage) -> Unit,
    ) {
        val scopeSessionId = module.approvalCoordinator.lookupSession(msg.requestId)
        if (scopeSessionId == null) {
            log.warn("handleApprovalResponse: no session found for requestId={}", msg.requestId.value)
            sendFrame(ServerMessage.ProtocolError("Unknown approval request"))
            return
        }
        module.approvalCoordinator.handleResponse(msg, scopeSessionId)?.let { sendFrame(it) }
    }

    // The message carries its own sessionId/stageId/requestId, so no lookup is needed — the parked
    // stage completes and re-runs with the answers in context (DefaultSessionOrchestrator).
    private suspend fun handleClarificationResponse(msg: ClientMessage.ClarificationResponse) {
        runCatching {
            module.orchestrator.submitClarification(msg.sessionId, msg.stageId, msg.requestId, msg.answers)
        }.onFailure {
            log.error("submitClarification failed for session={}: {}", msg.sessionId.value, it.message, it)
        }
    }

    // Tombstone an idea, landing the discard in the session that captured it (so all of one idea's
    // events stay co-located), then reply with the refreshed board. A missing idea is a no-op.
    private suspend fun handleDiscardIdea(
        msg: ClientMessage.DiscardIdea,
        sendFrame: suspend (ServerMessage) -> Unit,
    ) {
        runCatching {
            IdeaReader(module.eventStore).sessionOf(msg.ideaId)?.let { sessionId ->
                module.eventStore.append(
                    NewEvent(
                        metadata = EventMetadata(
                            eventId = EventId(UUID.randomUUID().toString()),
                            sessionId = sessionId,
                            timestamp = Clock.System.now(),
                            schemaVersion = 1,
                            causationId = null,
                            correlationId = null,
                        ),
                        payload = IdeaDiscardedEvent(ideaId = msg.ideaId, sessionId = sessionId),
                    ),
                )
            }
        }.onFailure {
            log.error("discardIdea failed for ideaId={}: {}", msg.ideaId, it.message, it)
        }
        sendFrame(queries.listIdeas())
    }

    private fun errorResponse(message: String) = ServerMessage.ProtocolError(message)

    private fun encodeError(message: String): Frame.Text =
        Frame.Text(ProtocolSerializer.encodeServerMessage(ServerMessage.ProtocolError(message)))

    private suspend fun handleCreateGrant(
        msg: ClientMessage.CreateGrant,
        sendFrame: suspend (ServerMessage) -> Unit,
    ) {
        // Validate: tiers must be non-empty and bounded to T2 (server-side ceiling).
        // T3/T4 are destructive/escalated tiers that must never be auto-approved via grants.
        if (msg.permittedTiers.isEmpty()) {
            sendFrame(errorResponse("CreateGrant: permittedTiers must not be empty"))
            return
        }
        val maxGrantableTier = Tier.T2
        val overBroad = msg.permittedTiers.filter { it.level > maxGrantableTier.level }
        if (overBroad.isNotEmpty()) {
            sendFrame(errorResponse("CreateGrant: tiers ${overBroad.map { it.name }} exceed maximum grantable tier ${maxGrantableTier.name}"))
            return
        }

        val scope = when (msg.scope) {
            GrantScopeDto.SESSION -> {
                if (msg.stageId != null) {
                    sendFrame(errorResponse("CreateGrant: SESSION scope must not include stageId"))
                    return
                }
                // Require toolName — blanket session grants (no operation scope) are rejected.
                val toolName = msg.toolName?.takeIf { it.isNotBlank() } ?: run {
                    sendFrame(errorResponse("CreateGrant: SESSION scope requires toolName to prevent blanket approval"))
                    return
                }
                GrantScope.SESSION(toolName)
            }
            GrantScopeDto.STAGE -> {
                val sid = msg.stageId ?: run {
                    sendFrame(errorResponse("CreateGrant: STAGE scope requires stageId"))
                    return
                }
                GrantScope.STAGE(sid)
            }
        }
        val event = NewEvent(
            metadata = EventMetadata(
                eventId = EventId(UUID.randomUUID().toString()),
                sessionId = msg.sessionId,
                timestamp = Clock.System.now(),
                schemaVersion = 1,
                causationId = null,
                correlationId = null,
            ),
            payload = ApprovalGrantCreatedEvent(
                grantId = GrantId(UUID.randomUUID().toString()),
                scope = scope,
                permittedTiers = msg.permittedTiers.toSet(),
                reason = msg.reason,
                expiresAt = msg.expiresAt,
                sessionId = msg.sessionId,
                stageId = (scope as? GrantScope.STAGE)?.stageId,
                projectId = null,
                toolName = (scope as? GrantScope.SESSION)?.toolName,
            ),
        )
        module.eventStore.append(event)
    }

    private suspend fun handleStartChatSession(
        session: DefaultWebSocketServerSession,
        msg: ClientMessage.StartChatSession,
        sendFrame: suspend (ServerMessage) -> Unit,
    ) {
        val sessionId = msg.sessionId
        log.info("starting chat session={}", sessionId.value)

        // Emit ChatSessionStartedEvent for the new session (no graph, no orchestrator).
        // streamGlobal maps it to a session.announced frame with the real sequence, and the
        // USER/ROUTER turns from onUserInput stream as chat.turn frames — all event-derived.
        module.eventStore.append(NewEvent(
            metadata = EventMetadata(
                eventId = EventId(UUID.randomUUID().toString()),
                sessionId = sessionId,
                timestamp = Clock.System.now(),
                schemaVersion = 1,
                causationId = null,
                correlationId = null,
            ),
            payload = ChatSessionStartedEvent(sessionId = sessionId),
        ))

        // Bind the per-repo project profile so router triage can cite conventions/commands.
        // Parity with launchSessionRun — chat sessions previously never bound a profile.
        runCatching { module.bindProjectProfile(sessionId) }
            .onFailure { log.warn("project profile binding failed for chat session={}: {}", sessionId.value, it.message) }

        withSessionContext(sessionId) {
            runCatching {
                module.routerFacade.onUserInput(sessionId, msg.text)
            }.onFailure {
                log.error("routerFacade.onUserInput failed: {}", it.message)
                sendFrame(ServerMessage.ProtocolError(
                    message = "Router error: ${it.message}",
                    sequence = null,
                    sessionSequence = null,
                ))
            }
        }
    }

    private suspend fun handleStartSession(
        session: DefaultWebSocketServerSession,
        msg: ClientMessage.StartSession,
        workingDir: String?,
    ) {
        val graph = module.workflowRegistry.find(msg.workflowId)
        log.info("find returned: {}", graph)
        if (graph == null) {
            val error = ServerMessage.ProtocolError(
                message = "Unknown workflowId: ${msg.workflowId}",
                sequence = null,
                sessionSequence = null,
            )
            session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(error)))
            return
        }
        val sessionId: SessionId = TypeId(UUID.randomUUID().toString())
        log.info("starting session={} workflow={}", sessionId.value, msg.workflowId)

        // Resolve workspace from the Hello-frame working directory (invariant #9: record the
        // resolver's decision at handshake time as an event; replay reads the recorded fact).
        val resolvedWorkspace: WorkspaceContext? = module.workspaceResolver?.let { resolver ->
            val resolution = withContext(Dispatchers.IO) { resolver.resolve(workingDir) }
            when (resolution) {
                is WorkspaceResolution.Bound -> {
                    log.info("workspace bound: root={}", resolution.workspace.workspaceRoot)
                    resolution.workspace
                }
                is WorkspaceResolution.Rejected -> {
                    // Deliberate: a rejected client path still binds the resolver's fallback
                    // workspace. The fallback root is recorded in SessionWorkspaceBoundEvent so
                    // replay reproduces the same binding; the rejection reason is audit-logged
                    // here and does NOT become a field on the event (YAGNI).
                    log.warn(
                        "workspace rejected: session={} reason={} fallback={}",
                        sessionId.value, resolution.reason, resolution.fallback.workspaceRoot,
                    )
                    resolution.fallback
                }
            }
        }

        // Emit SessionWorkspaceBoundEvent when a workspace was resolved (invariant #9).
        resolvedWorkspace?.let { ws ->
            module.eventStore.append(NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = Clock.System.now(),
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = SessionWorkspaceBoundEvent(
                    sessionId = sessionId,
                    workspaceRoot = ws.workspaceRoot.toString(),
                    allowedPaths = ws.allowedPaths.map { it.toString() },
                ),
            ))
        }

        // Send StageToolManifest FIRST. If the WS is already closed, this throws and the
        // orchestrator never launches — no silent data loss. The exception propagates to the
        // runCatching in handleClientMessage (caller) and is logged there.
        // No session announce is sent here — it is derived from WorkflowStartedEvent by
        // streamGlobal via DomainEventMapper (one session.announced per session with the
        // correct sequence number).
        val stages = graph.stages.map { (stageId, stageConfig) ->
            val tools = stageConfig.allowedTools.mapNotNull { toolName ->
                module.toolRegistry.resolve(toolName)?.let { tool ->
                    ToolDecl(name = tool.name, tier = tool.tier.level)
                }
            }
            StageToolDecl(stageId = stageId.value, tools = tools)
        }
        session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(
            ServerMessage.StageToolManifest(
                sessionId = sessionId,
                workflowId = msg.workflowId,
                stages = stages,
            ),
        )))

        // Thread the resolved workspace into the OrchestrationConfig so the live run uses
        // the bound workspace root for tool path containment and per-session tool registry.
        val sessionConfig = module.orchestrationConfig().copy(workspace = resolvedWorkspace)

        // Seed the freeform request (if any) before the run so it lands in the decision
        // journal and is pinned into every stage's context (the intent input channel).
        msg.input?.takeIf { it.isNotBlank() }?.let { intent ->
            module.eventStore.append(NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = Clock.System.now(),
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = InitialIntentEvent(sessionId = sessionId, intent = intent),
            ))
        }

        // Only launch orchestrator after protocol messages are delivered successfully.
        // Uses launchSessionRun() so the job is tracked in activeSessionJobs — this prevents
        // the OrchestrationResumedEvent subscription from spuriously launching a duplicate resume.
        module.launchSessionRun(sessionId, graph, sessionConfig)
    }
}

internal suspend fun streamGlobal(
    eventStore: EventStore,
    bridge: SessionEventBridge,
    mapper: DomainEventMapper,
    sendFrame: suspend (ServerMessage) -> Unit,
) = coroutineScope {
    val buffer = Channel<StoredEvent>(
        capacity = BUFFER_CAPACITY,
        onBufferOverflow = BufferOverflow.SUSPEND,
    )

    // Subscription must be active BEFORE replaySnapshot reads lastGlobalSequence,
    // otherwise events appended in that window are lost (replay=0 SharedFlow).
    val subscribed = CompletableDeferred<Unit>()
    val source = eventStore.subscribeAll()
    val signaled = when (source) {
        is SharedFlow<StoredEvent> -> source.onSubscription { subscribed.complete(Unit) }
        else -> source.onStart { subscribed.complete(Unit) }
    }
    val subscription = launch {
        signaled.collect { buffer.send(it) }
    }

    // Per-session sequence counters so live events carry the correct sessionSequence
    // and are not silently dropped by the TUI's SnapshotPhaseReducer dedup filter.
    val sessionSequences = mutableMapOf<String, Long>()

    try {
        subscribed.await()
        bridge.replaySnapshot()
        for (event in buffer) {
            val sid = event.metadata.sessionId.value
            val seq = sessionSequences.getOrDefault(sid, 0L) + 1
            sessionSequences[sid] = seq
            val msg = mapper.map(event, sessionSequence = seq) ?: continue
            sendFrame(msg)
        }
    } finally {
        subscription.cancel()
        buffer.close()
    }
}
