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 ->
            val kind = registry.get(s.produces)
                ?: throw WorkflowValidationException("Unknown artifact kind '${s.produces}' 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}'")
            }
            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)
        return WorkflowGraph(id = workflowId, stages = stageMap, transitions = transitions, start = start)
    }
}
