package com.correx.infrastructure.workflow

import com.correx.core.artifacts.kind.ArtifactKindRegistry
import com.correx.core.artifacts.kind.DefaultArtifactKindRegistry
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.StageConfig
import com.correx.core.transitions.graph.TransitionEdge
import com.correx.core.transitions.graph.WorkflowGraph
import com.fasterxml.jackson.annotation.JsonProperty
import com.fasterxml.jackson.databind.DeserializationFeature
import com.fasterxml.jackson.dataformat.toml.TomlMapper
import com.fasterxml.jackson.module.kotlin.kotlinModule
import com.fasterxml.jackson.module.kotlin.readValue
import java.nio.file.Path
import kotlin.io.path.exists
import kotlin.io.path.readText

private data class WorkflowFile(
    val id: String = "",
    val start: String = "",
    val description: String = "",
    val stages: List<StageSection> = emptyList(),
    val transitions: List<TransitionSection> = emptyList(),
)

data class WorkflowFileMeta(val id: String, val description: String)

private data class ProducesEntry(
    val name: String = "",
    val kind: String = "",
)

private data class StageSection(
    val id: String = "",
    val prompt: String? = null,
    @param:JsonProperty("system_prompt") val systemPrompt: String? = null,
    val produces: List<ProducesEntry> = emptyList(),
    val needs: List<String> = emptyList(),
    @param:JsonProperty("allowed_tools") val allowedTools: List<String> = emptyList(),
    val writes: List<String> = emptyList(),
    @param:JsonProperty("static_analysis") val staticAnalysis: List<String> = emptyList(),
    @param:JsonProperty("token_budget") val tokenBudget: Int = 4096,
    @param:JsonProperty("max_retries") val maxRetries: Int = 3,
    @param:JsonProperty("requires_approval") val requiresApproval: Boolean = false,
    @param:JsonProperty("inject_artifact_kinds") val injectArtifactKinds: Boolean = false,
    @param:JsonProperty("ground_references") val groundReferences: Boolean = false,
    @param:JsonProperty("brief_echo") val briefEcho: Boolean = false,
    @param:JsonProperty("require_task_decompose") val requireTaskDecompose: Boolean = false,
    @param:JsonProperty("claim_task") val claimTask: Boolean = false,
    @param:JsonProperty("semantic_review") val semanticReview: Boolean = false,
)

// condition fields flattened into the transition row
private data class TransitionSection(
    val id: String = "",
    val from: String = "",
    val to: String = "",
    @param:JsonProperty("condition_type") val conditionType: String = "",
    @param:JsonProperty("condition_artifact_id") val conditionArtifactId: String? = null,
    @param:JsonProperty("condition_key") val conditionKey: String? = null,
    @param:JsonProperty("condition_value") val conditionValue: String? = null,
    @param:JsonProperty("condition_field") val conditionField: String? = null,
    @param:JsonProperty("condition_operator") val conditionOperator: String? = null,
)

private const val TERMINAL = "done"

private val mapper = TomlMapper.builder()
    .addModule(kotlinModule())
    .disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
    .build()

class TomlWorkflowLoader(
    private val registry: ArtifactKindRegistry = DefaultArtifactKindRegistry(),
    // Operator sampling defaults applied to every stage's inference request (maxTokens is still pinned
    // per-stage to the token budget). Defaults reproduce the former hardcoded temperature=0.7/topP=1.0.
    private val samplingDefaults: GenerationConfig = GenerationConfig(temperature = 0.7, topP = 1.0, maxTokens = 0),
) : WorkflowLoader {
    override fun load(path: Path): WorkflowGraph {
        val raw = path.readText()
        val file = runCatching { mapper.readValue<WorkflowFile>(raw) }
            .getOrElse { throw WorkflowValidationException("Failed to parse workflow TOML: ${it.message}") }
        return file.toWorkflowGraph(path.parent)
    }

    override fun loadMeta(path: Path): WorkflowFileMeta {
        val raw = path.readText()
        val file = runCatching { mapper.readValue<WorkflowFile>(raw) }
            .getOrElse { throw WorkflowValidationException("Failed to parse workflow TOML: ${it.message}") }
        return WorkflowFileMeta(id = file.id, description = file.description)
    }

    private fun resolvePromptPath(raw: String, workflowDir: Path?): String {
        val p = Path.of(raw)
        if (p.isAbsolute) return raw
        val candidate = workflowDir?.resolve(raw)
        return if (candidate != null && candidate.exists()) candidate.toString() else raw
    }

    private fun WorkflowFile.toWorkflowGraph(workflowDir: Path?): WorkflowGraph {
        val startId = StageId(start)
        val stageMap = stages.associate { s ->
            StageId(s.id) to StageConfig(
                produces = s.produces.map { entry ->
                    val resolvedKind = registry.get(entry.kind)
                        ?: error("Unknown artifact kind '${entry.kind}' in stage '${s.id}'")
                    TypedArtifactSlot(name = ArtifactId(entry.name), kind = resolvedKind)
                },
                needs = s.needs.map { ArtifactId(it) }.toSet(),
                allowedTools = s.allowedTools.toSet(),
                writeManifest = s.writes,
                staticAnalysis = s.staticAnalysis,
                semanticReview = s.semanticReview,
                tokenBudget = s.tokenBudget,
                // Propagate the declared token budget to the inference completion cap.
                // Without this the StageConfig default (maxTokens=2048) is used, truncating
                // larger artifacts (finishReason=length) → invalid JSON → validation failure.
                generationConfig = samplingDefaults.copy(maxTokens = s.tokenBudget),
                maxRetries = s.maxRetries,
                metadata = s.toMetadata(workflowDir),
            )
        }

        val declaredIds = stageMap.keys.map { it.value }.toSet()
        validate(start, declaredIds, transitions)

        val allProduced = stageMap.values.flatMap { it.produces }.map { it.name }.toSet()
        stageMap.forEach { (stageId, config) ->
            config.needs.forEach { needed ->
                if (needed !in allProduced) {
                    throw WorkflowValidationException(
                        "Stage '${stageId.value}' needs artifact '${needed.value}' but no stage produces it"
                    )
                }
            }
        }

        val transitionSet = transitions.map { t ->
            TransitionEdge(
                id = TransitionId(t.id),
                from = StageId(t.from),
                to = StageId(t.to),
                condition = ConditionSpec(
                    type = t.conditionType,
                    artifactId = t.conditionArtifactId,
                    key = t.conditionKey,
                    value = t.conditionValue,
                    field = t.conditionField,
                    operator = t.conditionOperator,
                ).toCondition(),
            )
        }.toSet()

        return WorkflowGraph(id = id, stages = stageMap, transitions = transitionSet, start = startId)
    }

    private fun StageSection.toMetadata(workflowDir: Path?): Map<String, String> = buildMap {
        prompt?.let { put("prompt", resolvePromptPath(it, workflowDir)) }
        systemPrompt?.let { put("systemPrompt", resolvePromptPath(it, workflowDir)) }
        if (requiresApproval) put("requiresApproval", "true")
        if (injectArtifactKinds) put("injectArtifactKinds", "true")
        if (groundReferences) put("groundReferences", "true")
        if (briefEcho) put("briefEcho", "true")
        if (requireTaskDecompose) put("requireTaskDecompose", "true")
        if (claimTask) put("claimTask", "true")
    }

    @Suppress("ThrowsCount")
    private fun validate(
        start: String,
        declaredStages: Set<String>,
        transitions: List<TransitionSection>,
    ) {
        if (start !in declaredStages) {
            throw WorkflowValidationException("start stage '$start' is not declared in stages")
        }
        val seenIds = mutableSetOf<String>()
        transitions.forEach { t ->
            if (!seenIds.add(t.id)) {
                throw WorkflowValidationException("duplicate transition id: '${t.id}'")
            }
            checkTransitionStage(t, "from", t.from, declaredStages)
            checkTransitionStage(t, "to", t.to, declaredStages)
        }
    }

    private fun checkTransitionStage(t: TransitionSection, field: String, stageRef: String, declared: Set<String>) {
        if (stageRef !in declared && (field != "to" || stageRef != TERMINAL)) {
            throw WorkflowValidationException(
                "transition '${t.id}' references unknown $field-stage '$stageRef'"
            )
        }
    }
}
