package com.correx.core.kernel.orchestration

import com.correx.core.approvals.ApprovalOutcome
import com.correx.core.approvals.ApprovalStatus
import com.correx.core.approvals.model.ApprovalDecision
import com.correx.core.inference.Tokenizer
import com.correx.core.journal.DefaultDecisionJournalRepository
import com.correx.core.journal.DecisionJournalRenderer
import com.correx.core.artifacts.ArtifactState
import com.correx.core.artifacts.kind.ArtifactKindRegistry
import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.events.events.ApprovalDecisionResolvedEvent
import com.correx.core.events.events.ApprovalRequestedEvent
import com.correx.core.events.events.ArtifactContentStoredEvent
import com.correx.core.events.events.ClarificationAnswer
import com.correx.core.events.events.ClarificationAnsweredEvent
import com.correx.core.events.events.ArtifactValidatedEvent
import com.correx.core.events.events.OrchestrationResumedEvent
import com.correx.core.events.events.RefinementIterationEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.events.TransitionExecutedEvent
import com.correx.core.events.orchestration.OrchestrationState
import com.correx.core.events.types.ApprovalDecisionId
import com.correx.core.events.types.ApprovalRequestId
import com.correx.core.events.types.ArtifactId
import com.correx.core.events.types.ClarificationRequestId
import com.correx.core.events.types.ArtifactLifecyclePhase
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.subagent.InSessionSubagentRunner
import com.correx.core.kernel.orchestration.subagent.SubagentRunRequest
import com.correx.core.kernel.orchestration.subagent.SubagentRunner
import com.correx.core.kernel.retry.RetryCoordinator
import com.correx.core.sessions.Session
import com.correx.core.transitions.execution.StageExecutionResult
import com.correx.core.transitions.graph.WorkflowGraph
import com.correx.core.transitions.resolution.TransitionDecision
import kotlinx.datetime.Clock
import org.slf4j.LoggerFactory
import java.util.UUID
import java.util.concurrent.*
import java.util.concurrent.atomic.*

@SuppressWarnings(
    "ReturnCount",
    "NestedBlockDepth",
)
private val log = LoggerFactory.getLogger(DefaultSessionOrchestrator::class.java)

// Fallback refinement-loop cap when a stage does not declare maxRetries.
private const val DEFAULT_MAX_REFINEMENT = 3

class DefaultSessionOrchestrator(
    private val repositories: OrchestratorRepositories,
    engines: OrchestratorEngines,
    private val retryCoordinator: RetryCoordinator,
    artifactStore: ArtifactStore,
    tokenizer: Tokenizer? = null,
    private val decisionJournalRepository: DefaultDecisionJournalRepository,
    private val compactionService: JournalCompactionService? = null,
    artifactKindRegistry: ArtifactKindRegistry? = null,
    repoKnowledgeRetriever: RepoKnowledgeRetriever? = null,
) : SessionOrchestrator(repositories, engines, artifactStore, decisionJournalRepository, artifactKindRegistry = artifactKindRegistry, repoKnowledgeRetriever = repoKnowledgeRetriever), ApprovalGateway {
    override val tokenizer: Tokenizer? = tokenizer
    override val cancellations: ConcurrentHashMap<SessionId, AtomicBoolean> =
        ConcurrentHashMap<SessionId, AtomicBoolean>()

    override val subagentRunner: SubagentRunner = InSessionSubagentRunner(
        executeStage = { sid, stg, graph, session, cfg ->
            executeStage(sid, stg, graph, session, cfg, effectivesFor(cfg))
        },
    )

    override suspend fun run(
        sessionId: SessionId,
        graph: WorkflowGraph,
        config: OrchestrationConfig,
    ): WorkflowResult {
        log.debug("[Orchestrator] session={} workflow={} start={}", sessionId.value, graph.id, graph.start.value)
        emitWorkflowStarted(sessionId, graph, config)

        val base = ExecutionContext(graph, sessionId, 0, graph.start, config, null, null)
        val enriched = base.enrich()

        // Execute the start stage before entering the step loop
        return when (val result = enterStage(enriched, graph.start)) {
            is StepResult.Continue -> step(result.ctx)
            is StepResult.Terminal -> result.result
        }
    }

    @Suppress("LongMethod")
    private tailrec suspend fun step(ctx: EnrichedExecutionContext): WorkflowResult {
        log.debug(
            "[Orchestrator] step session={} stage={} stageCount={}",
            ctx.sessionId.value, ctx.currentStageId.value, ctx.stageCount,
        )
        if (isCancelled(ctx.sessionId)) return handleCancellation(ctx.sessionId, ctx.currentStageId)

        val enriched = EnrichedExecutionContext(
            graph = ctx.graph,
            sessionId = ctx.sessionId,
            stageCount = ctx.stageCount,
            currentStageId = ctx.currentStageId,
            config = ctx.config,
            state = orchestrationRepository.getState(ctx.sessionId),
            session = repositories.sessionRepository.getSession(ctx.sessionId),
        )

        val stageConfig = enriched.graph.stages[enriched.currentStageId]
        val targetIds: Set<ArtifactId> = stageConfig?.produces?.map { it.name }?.toSet() ?: emptySet()
        val stageArtifacts: Map<ArtifactId, ArtifactState> = if (targetIds.isEmpty()) {
            emptyMap()
        } else {
            val validated = repositories.eventStore.read(enriched.sessionId)
                .mapNotNull { (it.payload as? ArtifactValidatedEvent)?.artifactId }
                .filter { it in targetIds }
                .toSet()
            validated.associateWith { ArtifactState(phase = ArtifactLifecyclePhase.VALIDATED) }
        }

        val artifactContent: Map<ArtifactId, String> = targetIds.mapNotNull { id ->
            val cacheKey = "${enriched.sessionId.value}:${id.value}"
            artifactContentCache[cacheKey]?.let { content -> id to content }
        }.toMap()

        val resolved = resolveTransition(
            enriched.graph, enriched.sessionId, enriched.currentStageId, stageArtifacts, artifactContent,
        )
        // Preemptive redirect override (freestyle graph re-routing): if the operator confirmed a
        // jump (a recorded, unconsumed PreemptRedirectEvent), it overrides the resolver's edge —
        // unless the target is unknown or its needs aren't satisfied, in which case it is blocked
        // and the normal edge is taken. The decision is pure over the event log (replay-deterministic);
        // only the block-event emission is a side effect here.
        val decision = when (
            val outcome = PreemptRedirect.decide(
                events = repositories.eventStore.read(enriched.sessionId),
                graph = enriched.graph,
                sessionId = enriched.sessionId,
                artifactAvailable = { id ->
                    !artifactContentCache["${enriched.sessionId.value}:${id.value}"].isNullOrBlank()
                },
                nowMs = Clock.System.now().toEpochMilliseconds(),
            )
        ) {
            is PreemptRedirect.Outcome.Override -> outcome.move
            is PreemptRedirect.Outcome.Block -> {
                emit(enriched.sessionId, outcome.event)
                log.info(
                    "[Orchestrator] redirect blocked session={} to={} reason={}",
                    enriched.sessionId.value, outcome.event.toStageId.value, outcome.event.reason,
                )
                resolved
            }
            PreemptRedirect.Outcome.None -> resolved
        }
        log.debug(
            "[Orchestrator] transition session={} stage={} decision={}",
            enriched.sessionId.value, enriched.currentStageId.value, decision::class.simpleName,
        )
        return when (decision) {
            is TransitionDecision.Move -> when (val r = executeMove(enriched, decision)) {
                is StepResult.Continue -> step(r.ctx)
                is StepResult.Terminal -> r.result
            }

            is TransitionDecision.Stay -> if (isTerminal(enriched.graph, enriched.currentStageId)) {
                completeWorkflow(enriched.sessionId, enriched.currentStageId, enriched.stageCount)
            } else {
                failWorkflow(
                    sessionId = enriched.sessionId,
                    stageId = enriched.currentStageId,
                    reason = "no transition condition matched from stage ${enriched.currentStageId.value}",
                    retryExhausted = false,
                )
            }

            is TransitionDecision.Blocked -> failWorkflow(
                sessionId = enriched.sessionId,
                stageId = enriched.currentStageId,
                reason = decision.reason,
                retryExhausted = false,
            )

            is TransitionDecision.NoMatch -> failWorkflow(
                sessionId = enriched.sessionId,
                stageId = enriched.currentStageId,
                reason = "no matching transition from stage ${enriched.currentStageId.value}",
                retryExhausted = false,
            )
        }
    }


    /** Returns the cached validated artifact content for the given session + artifact id, or null if absent. */
    fun validatedArtifactContent(sessionId: SessionId, artifactId: ArtifactId): String? =
        artifactContentCache["${sessionId.value}:${artifactId.value}"]

    /**
     * Rebuilds [artifactContentCache] from durable events after a server restart.
     * Scans all events for [ArtifactContentStoredEvent], which records the
     * slot-name -> CAS content-hash mapping at write time, and re-fetches the
     * stored bytes from [artifactStore] by **content hash** (not slot name).
     * Must be called before [resume] so transition conditions and `needs`-injection
     * that read artifact content work correctly after a cold start.
     */
    suspend fun rehydrate(sessionId: SessionId) {
        repositories.eventStore.readFrom(sessionId, fromSequence = 0L).forEach { event ->
            val p = event.payload as? ArtifactContentStoredEvent ?: return@forEach
            val key = "${sessionId.value}:${p.artifactId.value}"
            artifactStore.get(p.contentHash)
                ?.toString(Charsets.UTF_8)
                ?.let { artifactContentCache[key] = it }
        }
    }

    suspend fun resume(
        sessionId: SessionId,
        graph: WorkflowGraph,
        config: OrchestrationConfig,
    ): WorkflowResult {
        val stageId = orchestrationRepository.getState(sessionId).currentStageId
            ?: return WorkflowResult.Failed(sessionId, "resume: no currentStageId", retryExhausted = false)
        if (!graph.stages.containsKey(stageId)) {
            // Terminal pseudo-stage ("done") or a stage from another graph — nothing to
            // re-enter. Treat as already completed instead of throwing inside enterStage.
            log.info(
                "[Orchestrator] resume: session={} stage={} not in graph '{}' — already terminal",
                sessionId.value, stageId.value, graph.id,
            )
            return WorkflowResult.Completed(sessionId, stageId)
        }
        log.info("[Orchestrator] resuming session={} workflow={} stage={}", sessionId.value, graph.id, stageId.value)
        val base = ExecutionContext(graph, sessionId, 0, stageId, config, null, null)
        val enriched = base.enrich()
        return when (val result = enterStage(enriched, stageId)) {
            is StepResult.Continue -> step(result.ctx.copy(currentStageId = stageId))
            is StepResult.Terminal -> result.result
        }
    }

    override suspend fun cancel(sessionId: SessionId) {
        cancellations.getOrPut(sessionId) { AtomicBoolean(false) }.set(true)
    }

    override suspend fun submitApprovalDecision(requestId: ApprovalRequestId, decision: ApprovalDecision) {
        val deferred = pendingApprovals[requestId]
        if (deferred != null) {
            deferred.complete(decision)
            return
        }
        // No live orchestrator coroutine for this request — server restarted while approval was pending.
        // Record the decision and emit resume to unblock the session in the event store.
        log.warn(
            "submitApprovalDecision: no live deferred for requestId={} (server restart?), emitting events directly",
            requestId.value,
        )
        val sessionId = decision.contextSnapshot.identity.sessionId
        val stageId = decision.contextSnapshot.identity.stageId
            ?: orchestrationRepository.getState(sessionId).currentStageId
            ?: return
        emit(sessionId, ApprovalDecisionResolvedEvent(
            decisionId = ApprovalDecisionId(UUID.randomUUID().toString()),
            requestId = requestId,
            outcome = decision.outcome ?: ApprovalOutcome.REJECTED,
            status = ApprovalStatus.COMPLETED,
            tier = decision.tier,
            resolutionTimestamp = Clock.System.now(),
            reason = decision.reason,
            userSteering = decision.userSteering,
        ))
        emit(sessionId, OrchestrationResumedEvent(sessionId, stageId))
    }

    /**
     * Delivers the operator's answers to a pending clarification raised by a stage. The waiting
     * stage (parked in [requestClarificationIfNeeded]) completes and re-runs with the answers in
     * context. If the server restarted while the clarification was pending there is no live
     * coroutine to complete, so the answers are recorded directly and the session resumed.
     */
    suspend fun submitClarification(
        sessionId: SessionId,
        stageId: StageId,
        requestId: ClarificationRequestId,
        answers: List<ClarificationAnswer>,
    ) {
        val deferred = pendingClarifications[requestId]
        if (deferred != null) {
            deferred.complete(answers)
            return
        }
        log.warn(
            "submitClarification: no live deferred for requestId={} (server restart?), recording directly",
            requestId.value,
        )
        emit(sessionId, ClarificationAnsweredEvent(requestId, sessionId, stageId, answers))
        emit(sessionId, OrchestrationResumedEvent(sessionId, stageId))
    }

    private suspend fun executeMove(
        ctx: EnrichedExecutionContext,
        decision: TransitionDecision.Move,
    ): StepResult {
        val nextStageId = decision.to

        if (isTerminal(ctx.graph, nextStageId)) {
            // Terminal stage is a sentinel — not executed, so don't count it.
            return StepResult.Terminal(
                completeWorkflow(ctx.sessionId, nextStageId, ctx.stageCount),
            )
        }

        // Runtime refinement guard: a back-edge (re-entering a stage already visited this
        // run, e.g. reviewer→implementer) increments a per-cycle counter recorded as an
        // event, so the guard is replay-deterministic. Exceeding the cap escalates to a
        // terminal failure instead of looping forever.
        if (isBackEdge(repositories.eventStore.read(ctx.sessionId), nextStageId)) {
            val cycleKey = "${ctx.currentStageId.value}->${nextStageId.value}"
            val maxIterations = ctx.graph.stages[nextStageId]?.maxRetries ?: DEFAULT_MAX_REFINEMENT
            val iteration = (orchestrationRepository.getState(ctx.sessionId).refinementIterations[cycleKey] ?: 0) + 1
            emit(ctx.sessionId, RefinementIterationEvent(ctx.sessionId, cycleKey, iteration, maxIterations))
            if (iteration > maxIterations) {
                return StepResult.Terminal(
                    failWorkflow(
                        ctx.sessionId,
                        nextStageId,
                        "refinement loop '$cycleKey' exceeded $maxIterations iterations — escalating",
                        retryExhausted = true,
                    ),
                )
            }
        }

        // Emit the transition before executing the next stage. Otherwise the next stage's
        // events (InferenceStarted, artifacts, …) are recorded ahead of the TransitionExecuted
        // that marks entering it, so the log shows the stage running before it was entered.
        val advancedTo = advanceStage(ctx.sessionId, ctx.currentStageId, decision)
        return when (val result = enterStage(ctx.copy(currentStageId = advancedTo), nextStageId)) {
            is StepResult.Continue -> StepResult.Continue(result.ctx)
            is StepResult.Terminal -> result
        }
    }

    @Suppress("ReturnCount")
    private suspend fun enterStage(
        ctx: EnrichedExecutionContext,
        stageId: StageId,
    ): StepResult {
        log.debug("[Orchestrator] executeStage session=${ctx.sessionId.value} stage=${stageId.value}")

        if (ctx.graph.stages[stageId]?.metadata?.get("requiresApproval") == "true") {
            val alreadyApproved = repositories.eventStore.read(ctx.sessionId).let { events ->
                val stageRequestIds = events
                    .mapNotNull { it.payload as? ApprovalRequestedEvent }
                    .filter { it.stageId == stageId && it.toolName == null }
                    .map { it.requestId }
                    .toSet()
                stageRequestIds.isNotEmpty() && events.any {
                    (it.payload as? ApprovalDecisionResolvedEvent)
                        ?.requestId in stageRequestIds
                }
            }
            if (!alreadyApproved) {
                val gateArtifact = ctx.graph.stages[stageId]?.needs?.firstOrNull()?.value
                val preview = gateArtifact
                    ?.let { artifactContentCache["${ctx.sessionId.value}:$it"] }
                    ?: "Upstream stage has completed. Review and approve to continue to stage ${stageId.value}."
                val approved = requestStageApproval(ctx.sessionId, stageId, preview)
                if (!approved) {
                    return StepResult.Terminal(
                        failWorkflow(ctx.sessionId, stageId, "approval rejected for stage ${stageId.value}", retryExhausted = false)
                    )
                }
            }
        }

        while (true) {
            if (isCancelled(ctx.sessionId)) {
                return StepResult.Terminal(handleCancellation(ctx.sessionId, stageId))
            }
            val result = subagentRunner.run(
                SubagentRunRequest(ctx.sessionId, stageId, ctx.graph, ctx.session, ctx.config),
            ).outcome
            when (result) {
                is StageExecutionResult.Success -> {
                    if (requestClarificationIfNeeded(ctx.sessionId, stageId, ctx.graph)) {
                        // The stage raised open questions and the operator answered them; loop to
                        // re-run the stage with the answers injected (no failure-retry budget spent).
                        continue
                    }
                    compactionService?.let { svc ->
                        val journalState = decisionJournalRepository.getJournal(ctx.sessionId)
                        val journalText = DecisionJournalRenderer().render(journalState)
                        val tokenEstimate = journalText.length / 4
                        svc.compactIfNeeded(
                            sessionId = ctx.sessionId,
                            state = journalState,
                            renderedTokenEstimate = tokenEstimate,
                            emit = { payload -> emit(ctx.sessionId, payload) },
                        )
                    }
                    return StepResult.Continue(ctx.copy(stageCount = ctx.stageCount + 1))
                }

                is StageExecutionResult.Failure -> {
                    log.debug(
                        "[Orchestrator] stage failed session=${ctx.sessionId.value} " +
                                "stage=${stageId.value} reason=${result.reason} retryable=${result.retryable}",
                    )
                    if (!result.retryable) {
                        // Non-retryable: let the transition resolver handle the outcome
                        // (e.g., back-edge via artifact_field_equals on failure)
                        return StepResult.Continue(ctx.copy(stageCount = ctx.stageCount + 1))
                    }
                    val refreshedState = orchestrationRepository.getState(ctx.sessionId)
                    val shouldRetry = retryCoordinator.shouldRetry(
                        sessionId = ctx.sessionId,
                        stageId = stageId,
                        currentAttempt = refreshedState.retryCount,
                        policy = ctx.config.retryPolicy,
                        failureReason = result.reason,
                    )
                    if (!shouldRetry) {
                        return StepResult.Terminal(
                            failWorkflow(
                                ctx.sessionId,
                                stageId,
                                result.reason,
                                retryExhausted = true,
                            ),
                        )
                    }
                    // retry — loop and re-execute
                }
            }
        }
    }

    private fun ExecutionContext.enrich() = EnrichedExecutionContext(
        graph, sessionId, stageCount, currentStageId, config,
        session = repositories.sessionRepository.getSession(sessionId),
        state = orchestrationRepository.getState(sessionId),
    )
}

// A back-edge re-enters a stage already transitioned into this run (e.g. reviewer→implementer).
// Top-level (not a member) to keep the orchestrator off the TooManyFunctions threshold.
private fun isBackEdge(events: List<StoredEvent>, target: StageId): Boolean =
    events.any { (it.payload as? TransitionExecutedEvent)?.to == target }

private sealed class StepResult {
    data class Continue(val ctx: EnrichedExecutionContext) : StepResult()
    data class Terminal(val result: WorkflowResult) : StepResult()
}

private data class ExecutionContext(
    val graph: WorkflowGraph,
    val sessionId: SessionId,
    val stageCount: Int,
    val currentStageId: StageId,
    val config: OrchestrationConfig,
    val session: Session?,
    val state: OrchestrationState?,
)

private data class EnrichedExecutionContext(
    val graph: WorkflowGraph,
    val sessionId: SessionId,
    val stageCount: Int,
    val currentStageId: StageId,
    val config: OrchestrationConfig,
    val session: Session,
    val state: OrchestrationState,
)
