package com.correx.infrastructure.workflow

import com.correx.core.artifacts.kind.ArtifactKindRegistry
import com.correx.core.artifacts.kind.TypedArtifactSlot
import com.correx.core.events.types.ArtifactId
import com.correx.core.events.types.StageId
import com.correx.core.events.types.TransitionId
import com.correx.core.inference.GenerationConfig
import com.correx.core.transitions.graph.BuildExpectation
import com.correx.core.transitions.graph.StageConfig
import com.correx.core.transitions.graph.TransitionEdge
import com.correx.core.transitions.graph.WorkflowGraph
import com.fasterxml.jackson.databind.DeserializationFeature
import com.fasterxml.jackson.databind.json.JsonMapper
import com.fasterxml.jackson.module.kotlin.kotlinModule
import com.fasterxml.jackson.module.kotlin.readValue
import java.nio.file.FileSystems
import java.nio.file.Path

private const val TERMINAL = "done"
private const val RECOVERY_STAGE = "recovery"
private const val DOD_ARTIFACT = "dod"
private const val STATIC_FLOOR_COMMAND = "correx-static-floor"

// Synthesized recovery stage prompt. Freestyle plans are LLM-emitted and never declare a recovery
// stage, but the kernel's retry-agency guard routes a write-less stage's unfixable gate failure to
// one when it exists. So the compiler always injects it (inert unless a gate actually routes): a
// write-capable stage that reads the failure ticket from context, repairs the cause, and hands
// control back to verification.
private const val RECOVERY_PROMPT =
    "You are the recovery stage. An upstream stage failed a verification gate it could not fix " +
        "(it lacked write access). The failure ticket — the failing stage, the gate, and its exact " +
        "output — is in your context under \"## Recovery ticket\". Read the gate output, inspect the " +
        "relevant files (file_read / list_dir), then FIX the reported cause. You MUST make the " +
        "change before finishing — investigation alone does not resolve the ticket, and handing " +
        "back without an edit only re-triggers the same failure.\n" +
        "Fix the actual cause the gate reports, whatever form it takes:\n" +
        "- Broken code / config → edit it with file_write.\n" +
        "- A MISSING file the build needs (e.g. package.json, tsconfig.json, a build/entry file " +
        "the gate can't find) → produce it. A build gate failing with \"no such file\" is fixed by " +
        "creating that file. If a whole project skeleton is missing, prefer the package manager's " +
        "own generator via shell — e.g. `npm create vite@latest <dir> -- --template react-ts` — " +
        "which writes a correct, complete file set (always pass the template/preset arg it " +
        "requires; the run is non-interactive so there's no prompt to fall back on). Fall back to " +
        "writing the files directly with file_write when no generator fits. `npm install` is " +
        "available to fetch deps once package.json exists. Only bare remote runners (`npx`, " +
        "`bunx`, `pnpx`) are blocked; use `npm create` instead.\n" +
        "Do not gratuitously rewrite code that already works — touch only what the gate blames. But " +
        "\"minimal\" means minimal to actually pass the gate, not zero.\n" +
        "When the cause is repaired, finish with a short resolution summary that mirrors the ticket: " +
        "the failing stage and gate, the root cause you found, and the exact files created/edited " +
        "and commands run. Then control returns to the stage that failed so its gate re-runs."

// Freestyle plans don't specify a per-stage token budget, so without this every compiled stage
// inherits StageConfig's 4096 default — 1/8th of what the static role_pipeline execution stages
// use (16384–32768). At 4096 a tool-heavy stage (scaffold/implement) truncates its context every
// round, evicting the file_read results the model just gathered, so it re-reads forever and never
// accumulates enough to write. Match the static execution baseline. (gemma ctx is 32768, so this
// leaves ample headroom for completion.)
private const val DEFAULT_STAGE_TOKEN_BUDGET = 24576

// A compiled freestyle stage must also lift its inference completion cap off StageConfig's 2048
// default, or the model is truncated (finishReason=length) mid-artifact — and a degenerating local
// model burns the whole 2048 on garbage (e.g. a `<|channel>thought` repetition loop) before it can
// stop. Mirror the static TomlWorkflowLoader path: pin the completion cap to the stage token budget.
class ExecutionPlanCompiler(
    private val registry: ArtifactKindRegistry,
    // Names of every registered tool. A stage that references a tool the runtime can't resolve
    // is dropped silently at run time (mapNotNull), so the stage runs toolless — invisible
    // without a live run. Validating here turns that into a compile-time rejection (which the
    // freestyle driver feeds back for a retry). Empty = the caller didn't supply the tool
    // universe, so validation is skipped (preserves the bare `ExecutionPlanCompiler(registry)`).
    private val knownTools: Set<String> = emptySet(),
    // When true, synthesize a write-capable `recovery` stage (+ a recovery->terminal edge) into every
    // compiled graph. The kernel's retry-agency guard routes a write-less stage's unfixable gate
    // failure there; without it the guard finds no recovery stage and falls back to futile in-place
    // retries (Vikunja #41). Off by default so the compiler's own unit tests see only plan stages;
    // the server's freestyle path (Main.kt) turns it on.
    private val injectRecovery: Boolean = false,
    // Operator sampling defaults for freestyle-compiled stages; maxTokens pinned to the stage budget.
    // Default reproduces the former hardcoded temperature=0.7/topP=1.0.
    private val samplingDefaults: GenerationConfig = GenerationConfig(temperature = 0.7, topP = 1.0, maxTokens = 0),
) {
    private val defaultStageGeneration = samplingDefaults.copy(maxTokens = DEFAULT_STAGE_TOKEN_BUDGET)
    private val mapper = JsonMapper.builder()
        .addModule(kotlinModule())
        .disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
        .build()

    /**
     * @param sessionArtifactIds artifact ids already produced earlier in the same freestyle session
     *   (e.g. the planning-phase `dod`). The execution plan never re-declares these, so they must be
     *   passed in for the cross-workflow `needs` seam (#264) to thread the DoD into impl/review stages.
     */
    fun compile(
        planJson: String,
        workflowId: String,
        sessionArtifactIds: Set<String> = emptySet(),
    ): WorkflowGraph {
        val plan = runCatching { mapper.readValue<ExecutionPlanModel>(planJson) }
            .getOrElse { throw WorkflowValidationException("Failed to parse execution_plan JSON: ${it.message}") }
        if (plan.stages.isEmpty()) throw WorkflowValidationException("execution_plan has no stages")
        validateTools(plan)
        validateScope(plan)

        // Parse + validate every stage's declared build_expectation up front — also feeds the
        // deterministic build-gate guarantee below.
        val declaredExpectations = plan.stages.associate { s ->
            s.id to (
                BuildExpectation.fromPlan(s.buildExpectation)
                    ?: throw WorkflowValidationException(
                        "stage '${s.id}' has invalid build_expectation '${s.buildExpectation}' — " +
                            "use one of none|module|project|tests",
                    )
                )
        }
        val autoGateStages = autoGateStages(plan, declaredExpectations)
        val sessionArtifacts =
            plan.stages.mapNotNull { it.produces.takeIf(String::isNotBlank) }.toSet() + sessionArtifactIds

        val stageMap = plan.stages.associate { s ->
            // `produces` is the unique slot name referenced by needs/edges; `kind` selects the
            // registered artifact kind. Plans predating the split used one string for both.
            val kindId = s.kind ?: s.produces
            val kind = registry.get(kindId)
                ?: throw WorkflowValidationException("Unknown artifact kind '$kindId' in stage '${s.id}'")
            // Write-guard manifest = static per-stage globs UNION the plan-derived set (BACKLOG
            // §B-§2). Plan stages carry no separate static `writes` glob list today, so the static
            // baseline here is empty and the plan-derived paths form the manifest; the union is kept
            // explicit so a future static source widens rather than replaces. Sorted for a
            // deterministic, replay-stable manifest.
            val writeManifest = PlanDerivedManifest.combine(
                staticManifest = emptyList(),
                planDerived = PlanDerivedManifest.deriveAllowedWrites(s),
            ).sorted()
            StageId(s.id) to StageConfig(
                produces = listOf(TypedArtifactSlot(name = ArtifactId(s.produces), kind = kind)),
                needs = (s.needs + requiredSessionArtifacts(s, sessionArtifacts)).map { ArtifactId(it) }.toSet(),
                // Every stage gets the full registered tool universe. The architect's per-stage
                // `tools` pick was advisory and mis-scoped stages (e.g. an edit-shaped stage granted
                // file_write but not file_edit → forced full-file rewrites; a repair stage unable to
                // touch a file it must fix). writeManifest still jails WHICH paths a stage may write,
                // so widening the tool grant does not widen write authority. Falls back to the
                // architect's pick only when no universe was supplied (bare-constructor unit tests).
                allowedTools = knownTools.ifEmpty { s.tools.toSet() },
                writeManifest = writeManifest,
                staticAnalysis = staticFloorCommands(writeManifest),
                // Option A: the concrete (non-glob) subset of the declared writes are expected files
                // the contract gate will require to exist. Globs (src/**) are write-permission scopes,
                // not existence guarantees, so they are excluded here.
                expectedFiles = s.writes.map { it.trim() }.filter { it.isNotEmpty() && !isGlob(it) },
                touches = s.touches.map { it.trim() }.filter { it.isNotEmpty() },
                buildExpectation = declaredExpectations.getValue(s.id),
                autoBuildGate = s.id in autoGateStages,
                semanticReview = s.semanticReview,
                tokenBudget = DEFAULT_STAGE_TOKEN_BUDGET,
                generationConfig = defaultStageGeneration,
                metadata = mapOf("role" to s.role, "promptInline" to reviewerBoundPrompt(s, DOD_ARTIFACT in sessionArtifacts)),
            )
        }
        // Inject a recovery stage unless the plan already declares one. Reached only by the kernel's
        // failure-ticket routing (no incoming plan edge), so it is excluded from the reachability
        // check below. At run time the kernel returns recovery to the exact stage that failed
        // (recoveryReturnMove, from the ticket origin); the recovery->terminal edge synthesized here
        // is only a fallback for the impossible case of a recovery stage reached with no ticket, so
        // the resolver always has a valid outgoing edge and can never dead-end.
        val terminal = terminalStageId(plan)
        val recoveryStage: Pair<StageId, StageConfig>? =
            if (!injectRecovery || RECOVERY_STAGE in stageMap.keys.map { it.value }) {
                null
            } else {
                StageId(RECOVERY_STAGE) to StageConfig(
                    allowedTools = knownTools.ifEmpty { setOf("file_write", "file_edit", "shell") },
                    tokenBudget = DEFAULT_STAGE_TOKEN_BUDGET,
                    generationConfig = defaultStageGeneration,
                    metadata = mapOf("role" to "recovery", "promptInline" to RECOVERY_PROMPT),
                )
            }
        val recoveryEdge: TransitionEdge? = recoveryStage?.let {
            TransitionEdge(
                id = TransitionId("$RECOVERY_STAGE->$terminal"),
                from = StageId(RECOVERY_STAGE),
                to = StageId(terminal),
                condition = ConditionSpec(type = "always_true").toCondition(),
            )
        }
        val fullStageMap = stageMap + listOfNotNull(recoveryStage)
        val declared = stageMap.keys.map { it.value }.toSet()

        val transitions = plan.edges.map { e ->
            if (e.from !in declared) {
                throw WorkflowValidationException("edge from unknown stage '${e.from}'")
            }
            if (e.to !in declared && e.to != TERMINAL) {
                throw WorkflowValidationException("edge to unknown stage '${e.to}'")
            }
            validateFieldCondition(plan, e)
            TransitionEdge(
                id = TransitionId("${e.from}->${e.to}"),
                from = StageId(e.from),
                to = StageId(e.to),
                condition = ConditionSpec(
                    type = e.condition.type,
                    artifactId = e.condition.artifactId,
                    key = e.condition.key,
                    field = e.condition.field,
                    value = e.condition.value,
                    operator = e.condition.operator,
                ).toCondition(),
            )
        }.toSet() + listOfNotNull(recoveryEdge)

        val start = StageId(plan.stages.first().id)
        // The recovery stage is reached by kernel routing, not a plan edge, so exclude it from the
        // "unreachable from start" check (it would otherwise always trip it).
        validateReachability(declared, transitions, start, excluded = setOfNotNull(recoveryStage?.first?.value))
        return WorkflowGraph(id = workflowId, stages = fullStageMap, transitions = transitions, start = start)
    }

    /**
     * Every tool a stage names must be a registered tool, else the runtime silently drops it
     * (resolve → null → mapNotNull) and the stage runs without it. Skipped when [knownTools] is
     * empty (the caller didn't supply the tool universe).
     */
    private fun validateTools(plan: ExecutionPlanModel) {
        if (knownTools.isEmpty()) return
        val offenders = plan.stages.flatMap { s -> s.tools.filter { it !in knownTools }.map { s.id to it } }
        if (offenders.isEmpty()) return
        val detail = offenders.joinToString(", ") { (stage, tool) -> "'$tool' in stage '$stage'" }
        throw WorkflowValidationException(
            "execution_plan references unknown tool(s): $detail — valid tools: " +
                knownTools.sorted().joinToString(", "),
        )
    }

    /**
     * Architecture-conformance scope check (design 2026-07-14). When a stage declares a `touches`
     * scope, every concrete (non-glob) path it also declares in `writes` must fall within it —
     * otherwise the stage has quietly widened past its stated intent (the observed failure: a
     * "web-approval-CLIENT" stage that also lists backend routes in its writes). Rejecting here, at
     * plan-compile time, hands the plan back to the architect via the existing plan-compile gate
     * BEFORE any code is written — the [writeManifest] guard cannot catch this because it is DERIVED
     * from the same `writes` and so can never contradict them. A glob in `writes` is a
     * write-permission scope, not a concrete target, so it is not checked against `touches` here.
     * A stage with no `touches` is unconstrained (opt-in), so existing plans are unaffected.
     */
    private fun validateScope(plan: ExecutionPlanModel) {
        for (s in plan.stages) {
            val touches = s.touches.map { it.trim() }.filter { it.isNotEmpty() }
            if (touches.isEmpty()) continue
            val matchers = touches.map { FileSystems.getDefault().getPathMatcher("glob:$it") }
            val escapees = s.writes
                .map { it.trim() }
                .filter { it.isNotEmpty() && !isGlob(it) }
                .filterNot { path -> matchers.any { it.matches(Path.of(path)) } }
            if (escapees.isNotEmpty()) {
                throw WorkflowValidationException(
                    "stage '${s.id}' declares scope touches=$touches but writes files outside it: " +
                        "${escapees.sorted()} — either keep the stage within its declared scope, or, " +
                        "if the change genuinely belongs to another area, widen `touches` or move " +
                        "those writes to a stage that owns that area.",
                )
            }
        }
    }

    /**
     * A field-equals edge can only ever fire if the producing stage's kind schema declares
     * that field — the kind's schema is also the LLM response format, so an undeclared field
     * (live repro: 'verdict' checked on an 'analysis' artifact) is never emitted and the
     * condition fails the whole workflow at runtime instead of the plan failing at lock time.
     */
    private fun validateFieldCondition(plan: ExecutionPlanModel, e: PlanEdge) {
        if (e.condition.type != "artifact_field_equals") return
        val field = e.condition.field ?: return
        val producing = plan.stages.firstOrNull { it.produces == e.condition.artifactId } ?: return
        val kindId = producing.kind ?: producing.produces
        val schema = registry.get(kindId)?.deriveJsonSchema() ?: return
        if (!schema.properties.containsKey(field)) {
            throw WorkflowValidationException(
                "edge '${e.from}->${e.to}' checks field '$field' of artifact '${e.condition.artifactId}', " +
                    "but its kind '$kindId' does not declare that field — use a kind whose schema has it " +
                    "(e.g. review_report for verdict)",
            )
        }
    }

    // The stage the plan ends on — the one whose edge points at "done" (a well-formed plan has exactly
    // one). Falls back to the last declared stage if no terminal edge is present yet.
    private fun terminalStageId(plan: ExecutionPlanModel): String =
        plan.edges.firstOrNull { it.to == TERMINAL }?.from ?: plan.stages.last().id

    private fun requiredSessionArtifacts(stage: PlanStage, sessionArtifacts: Set<String>): List<String> =
        if (
            DOD_ARTIFACT in sessionArtifacts &&
            stage.role.lowercase() in setOf("impl", "implementer", "review", "reviewer")
        ) {
            listOf(DOD_ARTIFACT)
        } else {
            emptyList()
        }

    private fun reviewerBoundPrompt(stage: PlanStage, dodPresent: Boolean): String {
        if (!dodPresent || stage.role.lowercase() !in setOf("review", "reviewer")) return stage.prompt
        return stage.prompt + "\n\n" +
            "The `dod` input artifact is your sole acceptance rubric. Judge only criteria tagged " +
            "`verified_by: reviewer`; deterministic gates own `verified_by: gate`. Approve if and only if " +
            "every reviewer criterion is met and no `out_of_scope` item was introduced. On rejection, " +
            "cite the unmet criterion ids. Never add or imply criteria outside the DoD."
    }

    // Deterministic build-gate floor. Code kinds carry an `imports_resolve` COMPILER-layer contract
    // assertion, but COMPILER assertions are enforced only by the execution gate, which fires only
    // when a stage sets build_expectation. The LLM planner emits every file-writing stage as the
    // generic `file_written` kind (never a typed code kind) and rarely sets build_expectation, so a
    // scaffold with a dangling import (`import './index.css'` for a file never written) sails through
    // to COMPLETED. Close it: unless a stage declares a real whole-project build (PROJECT/TESTS),
    // flag every write-declaring stage plus the terminal stage for an auto build gate — the terminal
    // one runs when the project is whole rather than failing legitimately-incomplete intermediate
    // stages. A MODULE declaration is only a per-file typecheck (delegated to LSP, which can silently
    // skip) and does NOT prove the assembled project builds, so it must not suppress the terminal
    // floor. Whether it ACTUALLY builds is decided at run time (SessionOrchestrator.runExecutionGate)
    // from the real FileWritten manifest — declared kinds don't reveal code-ness, written paths do —
    // so a docs-only plan is left alone and a code plan is always gated.
    private fun autoGateStages(
        plan: ExecutionPlanModel,
        declaredExpectations: Map<String, BuildExpectation>,
    ): Set<String> {
        val ownsRealBuild = declaredExpectations.values.any {
            it == BuildExpectation.PROJECT || it == BuildExpectation.TESTS
        }
        val writing = plan.stages.filter { it.writes.any { path -> path.isNotBlank() } }.map { it.id }.toSet()
        // A mid-plan PROJECT/TESTS declaration verifies nothing written after it, so it suppresses
        // the per-writing-stage gates (those builds would be redundant) but never the terminal floor.
        if (ownsRealBuild) return setOf(terminalStageId(plan))
        return if (writing.isEmpty()) emptySet() else writing + terminalStageId(plan)
    }

    private fun staticFloorCommands(writeManifest: List<String>): List<String> {
        return writeManifest.takeIf { it.isNotEmpty() }
            ?.let { listOf("$STATIC_FLOOR_COMMAND -- ${it.joinToString(" ")}") }
            ?: emptyList()
    }

    private fun isGlob(path: String): Boolean = path.any { it == '*' || it == '?' || it == '[' }

    private fun validateReachability(
        declared: Set<String>,
        transitions: Set<TransitionEdge>,
        start: StageId,
        excluded: Set<String> = emptySet(),
    ) {
        // A declared stage with no path from start can never run — typical LLM topology
        // failure: every stage wired straight to "done" instead of chained, which silently
        // truncates the plan to its first stage.
        val reachable = mutableSetOf(start.value)
        var frontier = setOf(start.value)
        while (frontier.isNotEmpty()) {
            frontier = transitions
                .filter { it.from.value in frontier }
                .map { it.to.value }
                .filter { it != TERMINAL && reachable.add(it) }
                .toSet()
        }
        val unreachable = declared - reachable - excluded
        if (unreachable.isNotEmpty()) {
            throw WorkflowValidationException(
                "stages unreachable from start '${start.value}': ${unreachable.sorted().joinToString(", ")} — " +
                    "chain stages with edges instead of pointing every stage to 'done'",
            )
        }
    }
}
