package com.correx.core.kernel.orchestration

import com.correx.core.approvals.domain.NoOpApprovalEngine
import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.journal.DefaultDecisionJournalRepository
import com.correx.core.journal.model.DecisionJournalState
import com.correx.core.sessions.projections.replay.EventReplayer
import com.correx.core.context.model.ContextPack
import com.correx.core.events.events.ArtifactValidatedEvent
import com.correx.core.events.events.ArtifactValidatingEvent
import com.correx.core.events.types.InferenceRequestId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.inference.GenerationConfig
import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.ResponseFormat
import com.correx.core.kernel.execution.ReplayStrategy
import com.correx.core.kernel.orchestration.subagent.InSessionSubagentRunner
import com.correx.core.kernel.orchestration.subagent.SubagentRunner
import com.correx.core.kernel.execution.WorkflowResult
import com.correx.core.kernel.replay.ReplayInferenceProvider
import com.correx.core.kernel.retry.exception.ReplayArtifactMissingException
import com.correx.core.risk.NoOpRiskAssessor
import com.correx.core.sessions.Session
import com.correx.core.transitions.execution.StageExecutionResult
import com.correx.core.transitions.graph.StageConfig
import com.correx.core.transitions.graph.WorkflowGraph
import com.correx.core.transitions.resolution.TransitionDecision
import com.correx.core.validation.model.ValidationContext
import org.slf4j.LoggerFactory
import java.nio.file.Path
import java.util.*
import java.util.concurrent.*
import java.util.concurrent.atomic.*

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

private data class ReplayContext(
    val graph: WorkflowGraph,
    val sessionId: SessionId,
    val currentStageId: StageId,
    val stageCount: Int,
    val config: OrchestrationConfig,
)

private sealed class ReplayStepResult {
    data class Continue(val ctx: ReplayContext) : ReplayStepResult()
    data class Terminal(val result: WorkflowResult) : ReplayStepResult()
}

/**
 * Replay is environment-independent. No live risk assessment or approval evaluation —
 * both are reconstructed from the event log (ApprovalGrantedEvent / ApprovalDeniedEvent).
 */
@SuppressWarnings("ReturnCount")
class ReplayOrchestrator(
    private val repositories: OrchestratorRepositories,
    engines: OrchestratorEngines,
    override val cancellations: ConcurrentHashMap<SessionId, AtomicBoolean>,
    private val strategy: ReplayStrategy,
    artifactStore: ArtifactStore,
) : SessionOrchestrator(
    repositories,
    engines.copy(approvalEngine = NoOpApprovalEngine(), riskAssessor = NoOpRiskAssessor()),
    artifactStore,
    DefaultDecisionJournalRepository(object : EventReplayer<DecisionJournalState> {
        override fun rebuild(sessionId: SessionId) = DecisionJournalState()
    }),
) {
    private val replayProvider = ReplayInferenceProvider(repositories.eventStore)

    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 {
        val session = repositories.sessionRepository.getSession(sessionId)

        // Derive workspace from the recorded SessionWorkspaceBoundEvent (invariants #1, #8, #9).
        // The resolver ran at handshake time; replay reads its recorded decision — no live FS calls.
        val effectiveConfig = session.state.boundWorkspace
            ?.let { bw ->
                config.copy(
                    workspace = WorkspaceContext(
                        workspaceRoot = Path.of(bw.workspaceRoot),
                        workingDir = Path.of(bw.workspaceRoot),
                        allowedPaths = bw.allowedPaths.map { Path.of(it) }.toSet(),
                    )
                )
            }
            ?: config

        emitWorkflowStarted(sessionId, graph, effectiveConfig)

        val ctx = ReplayContext(graph, sessionId, graph.start, 0, effectiveConfig)

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

    private tailrec suspend fun step(ctx: ReplayContext): WorkflowResult {
        if (isCancelled(ctx.sessionId)) return handleCancellation(ctx.sessionId, ctx.currentStageId)

        val session = repositories.sessionRepository.getSession(ctx.sessionId)

        return when (val decision = resolveTransition(ctx.graph, ctx.sessionId, ctx.currentStageId)) {
            is TransitionDecision.Move -> {
                val nextStageId = decision.to

                if (isTerminal(ctx.graph, nextStageId)) {
                    return completeWorkflow(ctx.sessionId, nextStageId, ctx.stageCount, ctx.graph.id)
                }

                when (val result = enterStage(ctx, nextStageId, session)) {
                    is ReplayStepResult.Continue -> step(
                        result.ctx.copy(
                            currentStageId = advanceStage(ctx.sessionId, ctx.currentStageId, decision),
                        ),
                    )
                    is ReplayStepResult.Terminal -> result.result
                }
            }

            is TransitionDecision.Stay -> if (isTerminal(ctx.graph, ctx.currentStageId)) {
                completeWorkflow(ctx.sessionId, ctx.currentStageId, ctx.stageCount, ctx.graph.id)
            } else {
                step(ctx)
            }

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

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

    private suspend fun enterStage(
        ctx: ReplayContext,
        stageId: StageId,
        session: Session,
    ): ReplayStepResult {
        log.debug("[Orchestrator] executeStage session=${ctx.sessionId.value} stage=${stageId.value}")
        val result =
            executeStage(ctx.sessionId, stageId, ctx.graph, session, ctx.config, effectivesFor(ctx.config))
        return when (result) {
            is StageExecutionResult.Success -> ReplayStepResult.Continue(
                ctx.copy(stageCount = ctx.stageCount + 1),
            )
            is StageExecutionResult.Failure ->
                ReplayStepResult.Terminal(
                    failWorkflow(ctx.sessionId, stageId, result.reason, retryExhausted = false),
                )
        }
    }

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

    @Suppress("LongParameterList")
    override suspend fun runInference(
        sessionId: SessionId,
        stageId: StageId,
        contextPack: ContextPack,
        stageConfig: StageConfig,
        timeoutMs: Long,
        responseFormat: ResponseFormat,
        effectives: RunEffectives,
        withTools: Boolean,
        forceWriteOnly: Boolean,
    ): InferenceResult = when (strategy) {
        is ReplayStrategy.SkipInference -> {
            // bypass router entirely — use recorded artifact
            val request = InferenceRequest(
                requestId = InferenceRequestId(UUID.randomUUID().toString()),
                sessionId = sessionId,
                stageId = stageId,
                contextPack = contextPack,
                generationConfig = GenerationConfig(
                    temperature = 0.0,
                    topP = 1.0,
                    maxTokens = 0,
                    seed = 0L,
                ),
                responseFormat = responseFormat,
            )
            runCatching { replayProvider.infer(request) }
                .fold(
                    onSuccess = { InferenceResult.Success(it) },
                    onFailure = { e ->
                        if (e is ReplayArtifactMissingException) {
                            InferenceResult.Failed(e.message ?: "replay artifact missing")
                        } else {
                            throw e
                        }
                    },
                )
        }
        else -> super.runInference(
            sessionId, stageId, contextPack, stageConfig, timeoutMs, responseFormat, effectives, withTools,
            forceWriteOnly,
        )
    }

    override suspend fun mapValidationOutcome(
        sessionId: SessionId,
        stageId: StageId,
        context: ValidationContext,
    ): StageExecutionResult = when (strategy) {
        is ReplayStrategy.SkipValidation -> {
            // Replay-only: emit a synthetic ArtifactValidatedEvent for every artifact left in
            // VALIDATING at this stage. Without this, artifacts are permanently stuck because
            // the real validation path (which emits ArtifactValidatedEvent) is bypassed.
            val events = repositories.eventStore.read(sessionId)
            val alreadyValidated = events
                .mapNotNull { (it.payload as? ArtifactValidatedEvent) }
                .filter { it.stageId == stageId }
                .map { it.artifactId }
                .toSet()
            events
                .mapNotNull { (it.payload as? ArtifactValidatingEvent)?.takeIf { e -> e.stageId == stageId } }
                .filter { it.artifactId !in alreadyValidated }
                .forEach { emit(sessionId, ArtifactValidatedEvent(it.artifactId, sessionId, stageId)) }
            StageExecutionResult.Success(emptyList())
        }
        else -> super.mapValidationOutcome(sessionId, stageId, context)
    }
}
