package com.correx.apps.server

import com.correx.apps.server.approval.ApprovalCoordinator
import com.correx.apps.server.freestyle.FreestyleDriver
import com.correx.apps.server.narration.NarrationSubscriber
import com.correx.apps.server.registry.ProviderRegistry
import com.correx.apps.server.registry.WorkflowRegistry
import com.correx.apps.server.workspace.WorkspaceResolver
import com.correx.core.approvals.ApprovalProjector
import com.correx.core.approvals.DefaultApprovalReducer
import com.correx.core.approvals.DefaultApprovalRepository
import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.events.events.ApprovalRequestedEvent
import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.OrchestrationResumedEvent
import com.correx.core.events.events.WorkflowFailedEvent
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.kernel.orchestration.DefaultSessionOrchestrator
import com.correx.core.kernel.orchestration.OrchestrationConfig
import com.correx.core.kernel.orchestration.OrchestrationRepository
import com.correx.core.router.RouterFacade
import com.correx.core.sessions.DefaultSessionRepository
import com.correx.core.sessions.SessionSummary
import com.correx.core.sessions.SessionSummaryProjector
import com.correx.core.tools.registry.ToolRegistry
import com.correx.core.events.orchestration.OrchestrationStatus
import kotlinx.datetime.Clock
import java.util.concurrent.ConcurrentHashMap
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.flow.filter
import kotlinx.coroutines.flow.launchIn
import kotlinx.coroutines.flow.onEach
import kotlinx.coroutines.launch
import org.slf4j.LoggerFactory

private val log = LoggerFactory.getLogger(ServerModule::class.java)

class ServerModule(
    val orchestrator: DefaultSessionOrchestrator,
    val eventStore: EventStore,
    val artifactStore: ArtifactStore,
    val sessionRepository: DefaultSessionRepository,
    val workflowRegistry: WorkflowRegistry,
    val providerRegistry: ProviderRegistry,
    val defaultOrchestrationConfig: OrchestrationConfig,
    val routerFacade: RouterFacade,
    val orchestrationRepository: OrchestrationRepository,
    val approvalRepository: DefaultApprovalRepository,
    val toolRegistry: ToolRegistry,
    val sessionUndoService: com.correx.apps.server.undo.SessionUndoService,
    // The managed-model router when correx owns the local model process ([[models]] configured);
    // null on the static-provider path. Backs the manual SwapModel / ClearModelPin operations.
    val modelSwapper: com.correx.infrastructure.inference.commons.ManagedInferenceRouter? = null,
    // Live GPU/RAM gauge pushed to clients on the global stream. Defaults to the always-unavailable
    // probe (static path / non-GPU host); the managed path wires an NVIDIA-backed probe.
    val resourceProbe: com.correx.infrastructure.inference.commons.ResourceProbe =
        com.correx.infrastructure.inference.commons.UnavailableProbe,
    // Resolves a client-supplied cwd into a trusted WorkspaceContext using the trust pipeline
    // (canonicalize → privileged check → allow-list clamp → boot fallback).
    val workspaceResolver: WorkspaceResolver? = null,
    // Long-lived scope owned by the module — backs the event-store subscription.
    // SupervisorJob so one failure doesn't kill the whole module;
    // Dispatchers.Default since the work is non-blocking and CPU-light.
    val moduleScope: CoroutineScope = CoroutineScope(SupervisorJob() + Dispatchers.Default),
    approvalCoordinator: ApprovalCoordinator? = null,
    val narrationMaxPerRun: Int = 100,
    // Cross-session, repo-scoped memory. Null disables the hooks entirely (tests / static path).
    val projectMemory: com.correx.apps.server.memory.ProjectMemoryService? = null,
    // Driver for the two-phase freestyle workflow: locks the plan and runs phase 2.
    // Null on non-freestyle paths; injected by Main when freestyle is configured.
    private val freestyleDriver: FreestyleDriver? = null,
) {
    val approvalCoordinator: ApprovalCoordinator = approvalCoordinator ?: ApprovalCoordinator(
        orchestrator = orchestrator,
    )

    private var subscriptionJob: Job? = null
    private val activeSessionJobs = ConcurrentHashMap<SessionId, Job>()

    /**
     * Subscribe to the event store and forward [ApprovalRequestedEvent] payloads to
     * the [ApprovalCoordinator]. Idempotent: calling twice is a no-op for the second call.
     *
     * Also pre-populates the coordinator's session lookup map from any approvals that were
     * pending when the server last stopped. This prevents a race where a reconnecting TUI
     * sends an ApprovalResponse before the per-connection replaySnapshot() has registered
     * the session mapping.
     */
    fun start() {
        if (subscriptionJob != null) return
        preRegisterPendingApprovals()
        resumeAbandonedSessions()
        subscriptionJob = eventStore.subscribeAll()
            .filter { it.payload is ApprovalRequestedEvent }
            .onEach { approvalCoordinator.onApprovalRequested(it.payload as ApprovalRequestedEvent) }
            .launchIn(moduleScope)

        // When a PAUSED session's approval is resolved without a live orchestrator coroutine
        // (server was restarted while approval was pending), emit OrchestrationResumedEvent
        // triggers this subscription to continue the session.
        eventStore.subscribeAll()
            .filter { it.payload is OrchestrationResumedEvent }
            .onEach { event ->
                val sessionId = (event.payload as OrchestrationResumedEvent).sessionId
                if (activeSessionJobs.containsKey(sessionId)) return@onEach
                val orchState = orchestrationRepository.getState(sessionId)
                if (orchState.status != OrchestrationStatus.RUNNING) return@onEach
                val graph = workflowRegistry.find(orchState.workflowId) ?: return@onEach
                launchSessionResume(sessionId, graph)
            }
            .launchIn(moduleScope)

        // Live-only: subscribeAll() replays nothing and ServerModule is never built under
        // ReplayOrchestrator, so narration never re-fires on restart/replay (invariant #8).
        NarrationSubscriber(eventStore = eventStore, routerFacade = routerFacade, scope = moduleScope, maxPerRun = narrationMaxPerRun).start()
    }

    /**
     * Launch the orchestrator's [run] for a new session and register the job in [activeSessionJobs].
     * This ensures the [OrchestrationResumedEvent] subscription guard
     * (`activeSessionJobs.containsKey(sessionId)`) works correctly — without it, the original
     * run job is invisible to the guard and a second `resume()` is spuriously launched,
     * duplicating inference calls (the root cause of the healthcheck script corruption).
     */
    fun launchSessionRun(
        sessionId: SessionId,
        graph: com.correx.core.transitions.graph.WorkflowGraph,
        sessionConfig: OrchestrationConfig = defaultOrchestrationConfig,
    ) {
        activeSessionJobs.computeIfAbsent(sessionId) {
            moduleScope.launch {
                // Record the repo map + seed prior-session memory before the run so stages
                // see both in context.
                projectMemory?.let { pm ->
                    runCatching {
                        pm.indexAndRecord(sessionId, pm.repoRoot())
                        pm.retrieveAndSeed(sessionId, pm.repoRoot())
                    }
                }
                runCatching {
                    orchestrator.run(sessionId, graph, sessionConfig)
                    // After a successful freestyle planning run, lock the plan and execute phase 2.
                    if (graph.id == "freestyle_planning") {
                        freestyleDriver?.lockAndRun(sessionId)
                    }
                    // Distil this run's decisions into durable project memory on completion.
                    projectMemory?.let { pm -> pm.persist(sessionId, pm.repoRoot()) }
                }.onFailure { ex ->
                    log.error("runSession: session={} failed: {}", sessionId.value, ex.message, ex)
                    eventStore.append(
                        NewEvent(
                            metadata = EventMetadata(
                                eventId = EventId(java.util.UUID.randomUUID().toString()),
                                sessionId = sessionId,
                                timestamp = Clock.System.now(),
                                schemaVersion = 1,
                                causationId = null,
                                correlationId = null,
                            ),
                            payload = WorkflowFailedEvent(
                                sessionId = sessionId,
                                stageId = graph.start,
                                reason = ex.message ?: "unexpected orchestrator failure",
                                retryExhausted = false,
                            ),
                        ),
                    )
                }
                activeSessionJobs.remove(sessionId)
            }
        }
    }

    private fun launchSessionResume(sessionId: SessionId, graph: com.correx.core.transitions.graph.WorkflowGraph) {
        activeSessionJobs.computeIfAbsent(sessionId) {
            moduleScope.launch {
                runCatching {
                    orchestrator.resume(sessionId, graph, defaultOrchestrationConfig)
                }.onFailure { e ->
                    log.error("resumeSession: session={} failed", sessionId.value, e)
                }
                activeSessionJobs.remove(sessionId)
            }
        }
    }

    /**
     * Rehydrates [artifactContentCache] from durable events and then resumes the session.
     * Called by [POST /sessions/{id}/resume] to restart a session after a server restart
     * without re-running already-completed stages.
     */
    fun launchSessionResumeWithRehydrate(sessionId: SessionId, graph: com.correx.core.transitions.graph.WorkflowGraph) {
        activeSessionJobs.computeIfAbsent(sessionId) {
            moduleScope.launch {
                runCatching {
                    orchestrator.rehydrate(sessionId)
                    orchestrator.resume(sessionId, graph, defaultOrchestrationConfig)
                }.onFailure { e ->
                    log.error("resumeSession (rehydrate): session={} failed", sessionId.value, e)
                }
                activeSessionJobs.remove(sessionId)
            }
        }
    }

    private fun resumeAbandonedSessions() {
        eventStore.allSessionIds().forEach { sessionId ->
            val orchState = orchestrationRepository.getState(sessionId)
            // Skip PAUSED sessions — they are waiting for user approval, not for computation.
            // preRegisterPendingApprovals() already routes their responses. When the user
            // approves, submitApprovalDecision emits OrchestrationResumedEvent which triggers
            // the subscription above to continue the session.
            if (orchState.status != OrchestrationStatus.RUNNING) return@forEach
            val graph = workflowRegistry.find(orchState.workflowId) ?: run {
                log.warn("resumeAbandoned: no graph for session={} workflowId={}", sessionId.value, orchState.workflowId)
                return@forEach
            }
            log.info("resumeAbandoned: launching session={} workflow={}", sessionId.value, orchState.workflowId)
            launchSessionResume(sessionId, graph)
        }
    }

    private val sessionSummaryProjector = SessionSummaryProjector()

    fun listSessionSummaries(): List<SessionSummary> =
        eventStore.allSessionIds().map { sessionId ->
            sessionSummaryProjector.project(sessionId, eventStore.readFrom(sessionId, fromSequence = 0L))
        }

    private fun preRegisterPendingApprovals() {
        val projector = ApprovalProjector(DefaultApprovalReducer())
        eventStore.allSessionIds().forEach { sessionId ->
            val orchState = orchestrationRepository.getState(sessionId)
            if (!orchState.pendingApproval) return@forEach
            val events = eventStore.read(sessionId)
            val approvalState = events.fold(projector.initial()) { state, event ->
                projector.apply(state, event)
            }
            approvalState.requests.values
                .filter { req -> approvalState.decisions.values.none { it.requestId == req.id } }
                .forEach { req -> approvalCoordinator.registerPendingRequest(req.id, sessionId) }
        }
    }
}
