package com.correx.apps.server

import com.correx.apps.server.approval.ApprovalCoordinator
import com.correx.apps.server.freestyle.FreestyleDriver
import com.correx.apps.server.logging.withSessionContext
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.config.OperatorProfile
import com.correx.core.config.ProjectProfileLoader
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.ExecutionPlanLockedEvent
import com.correx.core.events.events.OperatorProfileBoundEvent
import com.correx.core.events.events.ProjectProfileBoundEvent
import com.correx.core.events.events.OrchestrationPausedEvent
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.events.types.StageId
import com.correx.core.kernel.execution.WorkflowResult
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 kotlinx.datetime.DateTimeUnit
import kotlinx.datetime.minus
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 kotlinx.coroutines.withContext
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,
    // Rebuilt by ConfigService on a live config edit (router knobs); read at use-time by
    // GlobalStreamHandler.onUserInput so the next chat turn picks up the new facade.
    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).
    // Rebuilt by ConfigService when project.enabled is toggled live.
    projectMemory: com.correx.apps.server.memory.ProjectMemoryService? = null,
    // Live, swappable config. Null only in tests that don't exercise config editing; defaults to a
    // holder seeded from defaults so callers always have a value to read.
    val configHolder: com.correx.core.config.ConfigHolder =
        com.correx.core.config.ConfigHolder(com.correx.core.config.CorrexConfig()),
    // 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,
    // Operator profile snapshot loaded at server start. Null when personalization is disabled.
    private val operatorProfile: OperatorProfile? = null,
    // Propose learned profile updates at session end. Null when learn=false or no profile.
    // Rebuilt by ConfigService when personalization.enabled/learn is toggled live.
    profileAdaptationService: com.correx.apps.server.personalization.ProfileAdaptationService? = null,
    // Continuous health watch (observability-spec §4). Null disables it (tests / health.enabled=false);
    // records degraded/restored edge events on the system session, started in [start].
    private val healthMonitor: com.correx.apps.server.health.HealthMonitor? = null,
) {
    val approvalCoordinator: ApprovalCoordinator = approvalCoordinator ?: ApprovalCoordinator(
        orchestrator = orchestrator,
    )

    // Config-derived services that ConfigService swaps on a live edit. Reads happen at use-time
    // (next session / next chat turn), so in-flight sessions keep the config they started with.
    @Volatile
    var routerFacade: RouterFacade = routerFacade
        private set

    @Volatile
    var projectMemory: com.correx.apps.server.memory.ProjectMemoryService? = projectMemory
        private set

    @Volatile
    var profileAdaptationService: com.correx.apps.server.personalization.ProfileAdaptationService? =
        profileAdaptationService
        private set

    // Set by Main after construction (it needs a reference to the module to build the rebuild hook).
    lateinit var configService: com.correx.apps.server.config.ConfigService

    /**
     * Swaps the config-derived services after a live config edit. Called by [ConfigService].
     * In-flight sessions are untouched; the next session/turn reads the new instances.
     */
    fun applyRebuiltServices(
        newRouterFacade: RouterFacade,
        newProjectMemory: com.correx.apps.server.memory.ProjectMemoryService?,
        newProfileAdaptationService: com.correx.apps.server.personalization.ProfileAdaptationService?,
    ) {
        this.routerFacade = newRouterFacade
        this.projectMemory = newProjectMemory
        this.profileAdaptationService = newProfileAdaptationService
    }

    /** Orchestration config for a new session, with live-tunable knobs read from [configHolder]. */
    fun orchestrationConfig(): OrchestrationConfig = defaultOrchestrationConfig.copy(
        stageTimeoutMs = configHolder.get().orchestration.stageTimeoutMs,
    )

    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
                launchSessionResumeWithRehydrate(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()

        // Continuous, edge-triggered health watch (observability-spec §4). Live-only like narration:
        // probes read the environment and record degraded/restored events; replay reads those facts.
        healthMonitor?.start(moduleScope)
    }

    /**
     * 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 = orchestrationConfig(),
    ) {
        activeSessionJobs.computeIfAbsent(sessionId) {
            moduleScope.launch {
                withSessionContext(sessionId) {
                    // Record the repo map + seed prior-session memory before the run so stages
                    // see both in context.
                    projectMemory?.let { pm ->
                        runCatching {
                            pm.observeAndRecord(sessionId, pm.repoRoot())
                            pm.indexAndRecord(sessionId, pm.repoRoot())
                            pm.retrieveAndSeed(sessionId, pm.repoRoot())
                        }
                    }
                    // Bind operator profile snapshot as an event so replay reads the recorded
                    // fact rather than re-reading the live file (invariants #8/#9).
                    operatorProfile?.let { profile ->
                        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 = OperatorProfileBoundEvent(
                                    sessionId = sessionId,
                                    about = profile.about,
                                    approvalMode = profile.preferences.approvalMode,
                                    preferredModels = profile.preferences.preferredModels,
                                    conventions = profile.preferences.conventions,
                                ),
                            ),
                        )
                    }
                    bindProjectProfile(sessionId)
                    runCatching {
                        val result = orchestrator.run(sessionId, graph, sessionConfig)
                        freestyleHandoff(sessionId, graph, result)
                        // Distil this run's decisions into durable project memory on completion.
                        projectMemory?.let { pm -> pm.persist(sessionId, pm.repoRoot()) }
                        // Propose learned profile adaptations based on session journal (opt-in, never auto-applied).
                        operatorProfile?.let { profile ->
                            profileAdaptationService?.let { svc ->
                                runCatching { svc.proposeAdaptation(sessionId, profile) }
                                    .onFailure { log.warn("Profile adaptation failed: {}", it.message) }
                            }
                        }
                    }.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)
                }
            }
        }
    }

    /**
     * Resolves the workflow graph for resuming a freestyle phase-2 session. Phase-2 graphs
     * (`freestyle-<sessionId>`) are compiled from the locked plan, never registered in the
     * TOML registry, so the resume route cannot find them there. Rehydrates the artifact
     * cache first because the plan content lives in it after a restart.
     */
    suspend fun freestyleResumeGraph(sessionId: SessionId, workflowId: String): com.correx.core.transitions.graph.WorkflowGraph? {
        if (!workflowId.startsWith("freestyle-")) return null
        orchestrator.rehydrate(sessionId)
        return freestyleDriver?.compiledGraph(sessionId)
    }

    /**
     * Bind the per-repo project profile (curated .correx/project.toml) as an event so stages,
     * router chat triage, and replay read the recorded snapshot, never the live file
     * (invariants #8/#9). Shared by the workflow path and the chat-session path.
     */
    suspend fun bindProjectProfile(sessionId: SessionId) {
        val workspaceRoot = runCatching {
            sessionRepository.getSession(sessionId).state.boundWorkspace?.workspaceRoot
        }.getOrNull() ?: projectMemory?.repoRoot() ?: return
        val projectProfile = withContext(Dispatchers.IO) { ProjectProfileLoader.load(workspaceRoot) }
        if (projectProfile.isEmpty()) return
        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 = ProjectProfileBoundEvent(
                    sessionId = sessionId,
                    workspaceRoot = workspaceRoot,
                    about = projectProfile.about,
                    conventions = projectProfile.conventions,
                    commands = projectProfile.commands,
                ),
            ),
        )
    }

    /**
     * Phase-2 handoff for freestyle sessions, shared by every launcher that can finish a
     * planning graph (fresh run and all resume paths). Locks the plan and executes it,
     * guarded so a failed planning run never emits a spurious ExecutionPlanRejected and
     * a session whose plan is already locked never double-locks.
     */
    private suspend fun freestyleHandoff(
        sessionId: SessionId,
        graph: com.correx.core.transitions.graph.WorkflowGraph,
        result: WorkflowResult,
    ) {
        if (graph.id != "freestyle_planning") return
        if (result !is WorkflowResult.Completed) return
        val planAlreadyLocked = eventStore.read(sessionId)
            .any { it.payload is ExecutionPlanLockedEvent }
        if (planAlreadyLocked) return
        freestyleDriver?.lockAndRun(sessionId)
    }

    /**
     * Rehydrates [artifactContentCache] from durable events and then resumes the session.
     * Used by [POST /sessions/{id}/resume], boot-time abandoned-session pickup, and the
     * restart-while-approval-pending subscription — every resume needs the rehydrate
     * (pre-restart `needs`-artifacts live only in events/CAS) and the freestyle handoff
     * (a planning graph that completes on resume must still lock and run phase 2).
     */
    fun launchSessionResumeWithRehydrate(sessionId: SessionId, graph: com.correx.core.transitions.graph.WorkflowGraph) {
        activeSessionJobs.computeIfAbsent(sessionId) {
            moduleScope.launch {
                runCatching {
                    orchestrator.rehydrate(sessionId)
                    val result = orchestrator.resume(sessionId, graph, orchestrationConfig())
                    freestyleHandoff(sessionId, graph, result)
                }.onFailure { e ->
                    log.error("resumeSession (rehydrate): session={} failed", sessionId.value, e)
                }
                activeSessionJobs.remove(sessionId)
            }
        }
    }

    private fun resumeAbandonedSessions() {
        val maxAgeMinutes = configHolder.get().orchestration.resumeAbandonedMaxAgeMinutes
        if (maxAgeMinutes <= 0) {
            log.info("resumeAbandoned: disabled (resume_abandoned_max_age_minutes={})", maxAgeMinutes)
            return
        }
        val cutoff = Clock.System.now().minus(maxAgeMinutes, DateTimeUnit.MINUTE)
        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 lastEventAt = eventStore.read(sessionId).lastOrNull()?.metadata?.timestamp ?: return@forEach
            if (lastEventAt < cutoff) {
                log.info(
                    "resumeAbandoned: parking stale session={} (last event {}, cutoff {}m) — POST /resume revives it",
                    sessionId.value, lastEventAt, maxAgeMinutes,
                )
                // Record the parking decision as an event (invariant #9): the projection
                // otherwise reports RUNNING forever for a session no process is executing,
                // and the next boot would re-log the same skip.
                moduleScope.launch {
                    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 = OrchestrationPausedEvent(
                                sessionId = sessionId,
                                stageId = orchState.currentStageId ?: StageId.NONE,
                                reason = "ABANDONED_STALE",
                            ),
                        ),
                    )
                }
                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)
            launchSessionResumeWithRehydrate(sessionId, graph)
        }
    }

    private val sessionSummaryProjector = SessionSummaryProjector()

    fun listSessionSummaries(): List<SessionSummary> =
        eventStore.allSessionIds()
            // The system session carries global health events, not a user workflow — hide it.
            .filter { it != com.correx.apps.server.health.SYSTEM_SESSION }
            .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) }
        }
    }
}
