package com.correx.infrastructure.workflow

import com.correx.core.artifacts.kind.ConfigArtifactKind
import com.correx.core.artifacts.kind.DefaultArtifactKindRegistry
import com.correx.core.artifacts.kind.JsonSchema
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.assertThrows
import kotlin.test.assertEquals
import kotlin.test.assertNotNull

class ExecutionPlanCompilerTest {

    private val trivialSchema = JsonSchema(type = "object")
    private val registry = DefaultArtifactKindRegistry().also {
        it.register(ConfigArtifactKind(id = "patch", schema = trivialSchema, llmEmitted = true))
    }
    private val compiler = ExecutionPlanCompiler(registry)

    private val validPlan = """
        {
          "goal": "apply a patch",
          "stages": [
            {
              "id": "analyse",
              "prompt": "Analyse the codebase",
              "produces": "patch",
              "needs": [],
              "tools": ["ShellTool"]
            },
            {
              "id": "review",
              "prompt": "Review the patch",
              "produces": "patch",
              "needs": ["patch"],
              "tools": []
            }
          ],
          "edges": [
            {
              "from": "analyse",
              "to": "review",
              "condition": { "type": "always_true" }
            },
            {
              "from": "review",
              "to": "done",
              "condition": { "type": "always_true" }
            }
          ]
        }
    """.trimIndent()

    @Test
    fun `valid plan produces correct WorkflowGraph`() {
        val graph = compiler.compile(validPlan, "test-workflow")

        assertEquals("test-workflow", graph.id)
        assertEquals(2, graph.stages.size)
        assertEquals(2, graph.transitions.size)
        assertEquals("analyse", graph.start.value)

        val analyseStage = graph.stages[graph.start]
        assertNotNull(analyseStage)
        assertEquals("Analyse the codebase", analyseStage.metadata["promptInline"])
        assertEquals(setOf("patch"), analyseStage.produces.map { it.name.value }.toSet())
        assertEquals(setOf("ShellTool"), analyseStage.allowedTools)

        val reviewStage = graph.stages.values.first { it.metadata["promptInline"] == "Review the patch" }
        assertEquals(setOf("patch"), reviewStage.needs.map { it.value }.toSet())
    }

    @Test
    fun `edge referencing unknown from-stage throws WorkflowValidationException`() {
        val bad = validPlan.replace("\"from\": \"analyse\"", "\"from\": \"nonexistent\"")
        assertThrows<WorkflowValidationException> { compiler.compile(bad, "bad-workflow") }
    }

    @Test
    fun `edge referencing unknown to-stage throws WorkflowValidationException`() {
        val bad = validPlan.replace("\"to\": \"review\"", "\"to\": \"nonexistent\"")
        assertThrows<WorkflowValidationException> { compiler.compile(bad, "bad-workflow") }
    }

    @Test
    fun `unregistered produces kind throws WorkflowValidationException`() {
        val bad = validPlan.replace("\"produces\": \"patch\"", "\"produces\": \"unknown_kind\"")
        assertThrows<WorkflowValidationException> { compiler.compile(bad, "bad-workflow") }
    }

    @Test
    fun `empty stages throws WorkflowValidationException`() {
        val bad = """{ "goal": "x", "stages": [], "edges": [] }"""
        assertThrows<WorkflowValidationException> { compiler.compile(bad, "bad-workflow") }
    }

    @Test
    fun `malformed JSON throws WorkflowValidationException`() {
        assertThrows<WorkflowValidationException> { compiler.compile("not-json", "bad-workflow") }
    }
}
