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.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

private const val TERMINAL = "done"

// 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 = 16384

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(),
) {
    private val mapper = JsonMapper.builder()
        .addModule(kotlinModule())
        .disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
        .build()

    fun compile(planJson: String, workflowId: String): 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)

        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()
            val buildExpectation = BuildExpectation.fromPlan(s.buildExpectation)
                ?: throw WorkflowValidationException(
                    "stage '${s.id}' has invalid build_expectation '${s.buildExpectation}' — " +
                        "use one of none|module|project|tests",
                )
            StageId(s.id) to StageConfig(
                produces = listOf(TypedArtifactSlot(name = ArtifactId(s.produces), kind = kind)),
                needs = s.needs.map { ArtifactId(it) }.toSet(),
                allowedTools = s.tools.toSet(),
                writeManifest = 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) },
                buildExpectation = buildExpectation,
                semanticReview = s.semanticReview,
                tokenBudget = DEFAULT_STAGE_TOKEN_BUDGET,
                metadata = mapOf("promptInline" to s.prompt),
            )
        }
        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()

        val start = StageId(plan.stages.first().id)
        validateReachability(declared, transitions, start)
        return WorkflowGraph(id = workflowId, stages = stageMap, 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(", "),
        )
    }

    /**
     * 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)",
            )
        }
    }

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

    private fun validateReachability(declared: Set<String>, transitions: Set<TransitionEdge>, start: StageId) {
        // 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
        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'",
            )
        }
    }
}
