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

class ExecutionPlanCompiler(
    private val registry: ArtifactKindRegistry,
) {
    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")

        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}'")
            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(),
                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)
    }

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