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 com.correx.core.artifacts.kind.JsonSchemaProperty
import com.correx.core.events.types.StageId
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.assertThrows
import kotlin.test.assertEquals
import kotlin.test.assertFalse
import kotlin.test.assertNotNull
import kotlin.test.assertTrue

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 recoveringCompiler = ExecutionPlanCompiler(registry, injectRecovery = true)

    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 `injectRecovery synthesizes a write-capable recovery stage looping back to the terminal`() {
        val graph = recoveringCompiler.compile(validPlan, "recovery-workflow")

        // Two plan stages + the synthesized recovery stage; still reachable-valid (recovery is
        // excluded from the reachability check because it's entered by kernel routing, not an edge).
        assertEquals(3, graph.stages.size)
        val recovery = graph.stages[StageId("recovery")]
        assertNotNull(recovery)
        assertEquals("recovery", recovery.metadata["role"])
        assertTrue(recovery.allowedTools.contains("file_write"))
        // file_read / list_dir are always available to any tool-granting stage.
        assertTrue(recovery.effectiveAllowedTools.contains("file_read"))

        // recovery -> review (the plan's terminal stage, the one whose edge points at 'done').
        val recoveryEdge = graph.transitions.single { it.from == StageId("recovery") }
        assertEquals(StageId("review"), recoveryEdge.to)
    }

    @Test
    fun `injectRecovery is a no-op when the plan already declares a recovery stage`() {
        val plan = """
            {
              "goal": "x",
              "stages": [
                { "id": "recovery", "prompt": "fix", "produces": "patch", "needs": [], "tools": ["file_write"] }
              ],
              "edges": [
                { "from": "recovery", "to": "done", "condition": { "type": "always_true" } }
              ]
            }
        """.trimIndent()
        val graph = recoveringCompiler.compile(plan, "wf")
        assertEquals(1, graph.stages.size)
    }

    @Test
    fun `plan-declared writes are derived into the stage write manifest`() {
        val plan = """
            {
              "goal": "implement a feature",
              "stages": [
                {
                  "id": "impl",
                  "prompt": "Implement it",
                  "produces": "patch",
                  "needs": [],
                  "tools": ["file_write"],
                  "writes": ["core/feature/**", "docs/feature.md"]
                }
              ],
              "edges": [
                { "from": "impl", "to": "done", "condition": { "type": "always_true" } }
              ]
            }
        """.trimIndent()
        val graph = compiler.compile(plan, "manifest-workflow")
        val stage = graph.stages.values.single()
        assertEquals(setOf("core/feature/**", "docs/feature.md"), stage.writeManifest.toSet())
    }

    @Test
    fun `a stage without declared writes has an empty write manifest`() {
        val graph = compiler.compile(validPlan, "no-writes-workflow")
        assertTrue(graph.stages.values.all { it.writeManifest.isEmpty() })
    }

    @Test
    fun `writes escaping declared touches scope throws WorkflowValidationException naming the escapee`() {
        // The observed failure: a "client" stage confined to frontend that also lists a backend file.
        val plan = """
            {
              "goal": "build the approval client",
              "stages": [
                {
                  "id": "impl_client",
                  "prompt": "Build the web approval client",
                  "produces": "patch",
                  "needs": [],
                  "tools": ["file_write"],
                  "touches": ["frontend/**"],
                  "writes": ["frontend/src/pages/Approvals.tsx", "apps/server/routes/ApprovalStageRoute.kt"]
                }
              ],
              "edges": [
                { "from": "impl_client", "to": "done", "condition": { "type": "always_true" } }
              ]
            }
        """.trimIndent()
        val ex = assertThrows<WorkflowValidationException> { compiler.compile(plan, "scope-workflow") }
        assertTrue(
            ex.message!!.contains("apps/server/routes/ApprovalStageRoute.kt"),
            "the rejection names the out-of-scope write: ${ex.message}",
        )
    }

    @Test
    fun `writes within declared touches scope compile, and globs are not scope-checked`() {
        val plan = """
            {
              "goal": "build the approval client",
              "stages": [
                {
                  "id": "impl_client",
                  "prompt": "Build the web approval client",
                  "produces": "patch",
                  "needs": [],
                  "tools": ["file_write"],
                  "touches": ["frontend/**"],
                  "writes": ["frontend/src/pages/Approvals.tsx", "frontend/src/**"]
                }
              ],
              "edges": [
                { "from": "impl_client", "to": "done", "condition": { "type": "always_true" } }
              ]
            }
        """.trimIndent()
        val graph = compiler.compile(plan, "scope-ok-workflow")
        assertEquals(listOf("frontend/**"), graph.stages.getValue(StageId("impl_client")).touches)
    }

    @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 `kind field selects the registered kind while produces names the slot`() {
        val plan = """
            {
              "goal": "create files",
              "stages": [
                {
                  "id": "create_files",
                  "prompt": "Create the files",
                  "produces": "files_created",
                  "kind": "file_written",
                  "needs": [],
                  "tools": ["file_write"]
                }
              ],
              "edges": [
                { "from": "create_files", "to": "done", "condition": { "type": "always_true" } }
              ]
            }
        """.trimIndent()
        val graph = compiler.compile(plan, "kind-workflow")
        val slot = graph.stages.values.single().produces.single()
        assertEquals("files_created", slot.name.value)
        assertEquals("file_written", slot.kind.id)
    }

    @Test
    fun `unregistered kind field throws WorkflowValidationException`() {
        val plan = """
            {
              "goal": "x",
              "stages": [
                { "id": "s1", "prompt": "p", "produces": "out", "kind": "made_up_kind" }
              ],
              "edges": [
                { "from": "s1", "to": "done", "condition": { "type": "always_true" } }
              ]
            }
        """.trimIndent()
        assertThrows<WorkflowValidationException> { compiler.compile(plan, "bad-workflow") }
    }

    @Test
    fun `plan with unreachable stages throws WorkflowValidationException`() {
        // Live repro (2026-06-12): the architect wired every stage straight to "done"
        // instead of chaining, so the workflow completed after stage 1 of 3.
        val plan = """
            {
              "goal": "create three files",
              "stages": [
                { "id": "create_1", "prompt": "p1", "produces": "f1", "kind": "file_written" },
                { "id": "create_2", "prompt": "p2", "produces": "f2", "kind": "file_written" },
                { "id": "create_3", "prompt": "p3", "produces": "f3", "kind": "file_written" }
              ],
              "edges": [
                { "from": "create_1", "to": "done", "condition": { "type": "artifact_validated", "artifact_id": "f1" } },
                { "from": "create_2", "to": "done", "condition": { "type": "artifact_validated", "artifact_id": "f2" } },
                { "from": "create_3", "to": "done", "condition": { "type": "artifact_validated", "artifact_id": "f3" } }
              ]
            }
        """.trimIndent()
        val ex = assertThrows<WorkflowValidationException> { compiler.compile(plan, "bad-workflow") }
        assertEquals(true, ex.message?.contains("unreachable"))
    }

    @Test
    fun `chained plan with single terminal edge compiles`() {
        val plan = """
            {
              "goal": "create three files",
              "stages": [
                { "id": "create_1", "prompt": "p1", "produces": "f1", "kind": "file_written" },
                { "id": "create_2", "prompt": "p2", "produces": "f2", "kind": "file_written" },
                { "id": "create_3", "prompt": "p3", "produces": "f3", "kind": "file_written" }
              ],
              "edges": [
                { "from": "create_1", "to": "create_2", "condition": { "type": "artifact_validated", "artifact_id": "f1" } },
                { "from": "create_2", "to": "create_3", "condition": { "type": "artifact_validated", "artifact_id": "f2" } },
                { "from": "create_3", "to": "done", "condition": { "type": "artifact_validated", "artifact_id": "f3" } }
              ]
            }
        """.trimIndent()
        val graph = compiler.compile(plan, "chained-workflow")
        assertEquals(3, graph.stages.size)
        assertEquals("create_1", graph.start.value)
    }

    @Test
    fun `field condition on a kind without that field throws WorkflowValidationException`() {
        // Live repro (2026-06-12): a reviewer stage with kind=analysis and a verdict-equals
        // edge — analysis never emits 'verdict', so the workflow failed at runtime.
        val plan = """
            {
              "goal": "verify",
              "stages": [
                { "id": "verify", "prompt": "verify", "produces": "verify_out", "kind": "patch" }
              ],
              "edges": [
                {
                  "from": "verify", "to": "done",
                  "condition": { "type": "artifact_field_equals", "artifact_id": "verify_out", "field": "verdict", "value": "approved", "operator": "eq" }
                }
              ]
            }
        """.trimIndent()
        val ex = assertThrows<WorkflowValidationException> { compiler.compile(plan, "bad-workflow") }
        assertEquals(true, ex.message?.contains("verdict"))
    }

    @Test
    fun `field condition on a kind declaring the field compiles`() {
        val verdictRegistry = DefaultArtifactKindRegistry().also {
            it.register(
                ConfigArtifactKind(
                    id = "review_report",
                    schema = JsonSchema(
                        type = "object",
                        properties = mapOf("verdict" to JsonSchemaProperty(type = "string")),
                        required = listOf("verdict"),
                        additionalProperties = true,
                    ),
                    llmEmitted = true,
                ),
            )
        }
        val plan = """
            {
              "goal": "verify",
              "stages": [
                { "id": "verify", "prompt": "verify", "produces": "verify_out", "kind": "review_report" }
              ],
              "edges": [
                {
                  "from": "verify", "to": "done",
                  "condition": { "type": "artifact_field_equals", "artifact_id": "verify_out", "field": "verdict", "value": "approved", "operator": "eq" }
                }
              ]
            }
        """.trimIndent()
        val graph = ExecutionPlanCompiler(verdictRegistry).compile(plan, "verdict-workflow")
        assertEquals(1, graph.stages.size)
    }

    @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") }
    }

    // A compiler that knows the tool universe, so plan tool names are validated.
    private val toolAware = ExecutionPlanCompiler(
        registry,
        knownTools = setOf("ShellTool", "file_write", "task_context", "task_update"),
    )

    @Test
    fun `unknown tool name throws WorkflowValidationException naming the offender`() {
        // The runtime would silently drop a misspelled tool; here it fails to compile instead.
        val bad = validPlan.replace("\"tools\": [\"ShellTool\"]", "\"tools\": [\"task_updates\"]")
        val ex = assertThrows<WorkflowValidationException> { toolAware.compile(bad, "bad-workflow") }
        assertEquals(true, ex.message?.contains("task_updates"))
        assertEquals(true, ex.message?.contains("analyse"))
    }

    @Test
    fun `plan threading registered task tools compiles`() {
        val plan = """
            {
              "goal": "implement and review a tracked task",
              "stages": [
                {
                  "id": "implement", "prompt": "claim auth-1, build, submit", "produces": "patch",
                  "needs": [], "tools": ["file_write", "task_context", "task_update"]
                },
                {
                  "id": "review", "prompt": "review and complete auth-1", "produces": "patch",
                  "needs": ["patch"], "tools": ["task_context", "task_update"]
                }
              ],
              "edges": [
                { "from": "implement", "to": "review", "condition": { "type": "always_true" } },
                { "from": "review", "to": "done", "condition": { "type": "always_true" } }
              ]
            }
        """.trimIndent()
        val graph = toolAware.compile(plan, "tracked-workflow")
        val implement = graph.stages.values.first { it.metadata["promptInline"] == "claim auth-1, build, submit" }
        assertTrue(implement.allowedTools.containsAll(setOf("task_context", "task_update")))
    }

    @Test
    fun `tool validation is skipped when the tool universe is unknown`() {
        // The default compiler has no knownTools, so an arbitrary tool name still compiles.
        val plan = validPlan.replace("\"tools\": [\"ShellTool\"]", "\"tools\": [\"anything_goes\"]")
        val graph = compiler.compile(plan, "unvalidated-workflow")
        assertEquals(2, graph.stages.size)
    }

    @Test
    fun `Option A - concrete declared writes become expectedFiles, globs excluded`() {
        val plan = """
            {
              "goal": "x", "stages": [
                { "id": "impl", "prompt": "p", "produces": "patch", "tools": ["file_write"],
                  "writes": ["src/App.tsx", "src/hooks/**", "package.json"] }
              ],
              "edges": [ { "from": "impl", "to": "done", "condition": { "type": "always_true" } } ]
            }
        """.trimIndent()
        val stage = compiler.compile(plan, "wf").stages.values.single()
        assertEquals(setOf("src/App.tsx", "package.json"), stage.expectedFiles.toSet())
    }

    @Test
    fun `build_expectation is parsed into the closed-set enum`() {
        val plan = validPlan.replace(
            "\"tools\": [\"ShellTool\"]",
            "\"tools\": [\"ShellTool\"], \"build_expectation\": \"project\"",
        )
        val stage = compiler.compile(plan, "wf").stages[compiler.compile(plan, "wf").start]!!
        assertEquals(com.correx.core.transitions.graph.BuildExpectation.PROJECT, stage.buildExpectation)
    }

    @Test
    fun `absent build_expectation defaults to NONE`() {
        val stage = compiler.compile(validPlan, "wf").stages.values.first()
        assertEquals(com.correx.core.transitions.graph.BuildExpectation.NONE, stage.buildExpectation)
    }

    @Test
    fun `invalid build_expectation is rejected at compile time`() {
        val bad = validPlan.replace(
            "\"tools\": [\"ShellTool\"]",
            "\"tools\": [\"ShellTool\"], \"build_expectation\": \"deploy\"",
        )
        assertThrows<WorkflowValidationException> { compiler.compile(bad, "wf") }
    }

    // --- deterministic build-gate guarantee -------------------------------------------------------

    // Registry whose kinds resolve to code contracts (react_component / react_entry carry the
    // `imports_resolve` COMPILER assertion in KindContractTable).
    private val codeRegistry = DefaultArtifactKindRegistry().also {
        it.register(ConfigArtifactKind(id = "react_component", schema = trivialSchema, llmEmitted = false))
        it.register(ConfigArtifactKind(id = "react_entry", schema = trivialSchema, llmEmitted = false))
    }
    private val codeCompiler = ExecutionPlanCompiler(codeRegistry)

    private val codePlanNoGate = """
        {
          "goal": "scaffold a react app",
          "stages": [
            { "id": "entry", "prompt": "write main.tsx", "produces": "app_entry", "kind": "react_entry",
              "needs": [], "tools": ["file_write"], "writes": ["src/main.tsx"] },
            { "id": "views", "prompt": "write components", "produces": "app_views", "kind": "react_component",
              "needs": ["app_entry"], "tools": ["file_write"], "writes": ["src/App.tsx"] }
          ],
          "edges": [
            { "from": "entry", "to": "views", "condition": { "type": "always_true" } },
            { "from": "views", "to": "done", "condition": { "type": "always_true" } }
          ]
        }
    """.trimIndent()

    @Test
    fun `every write-declaring stage of an ungated plan is flagged for an auto build gate`() {
        val graph = codeCompiler.compile(codePlanNoGate, "wf")
        val views = graph.stages[StageId("views")]!!
        assertTrue(
            views.autoBuildGate,
            "a write-declaring stage of an ungated plan must be flagged for an auto build gate",
        )
        assertEquals(
            com.correx.core.transitions.graph.BuildExpectation.NONE,
            views.buildExpectation,
            "the flag alone — declared expectation stays NONE; the runtime promotes to PROJECT only if a " +
                "code module was actually written",
        )
        assertTrue(
            graph.stages[StageId("entry")]!!.autoBuildGate,
            "implementation stages must not defer compiler coverage to a terminal verifier",
        )
        assertTrue(graph.stages[StageId("entry")]!!.staticAnalysis.single().contains("src/main.tsx"))
        assertTrue(graph.stages[StageId("views")]!!.staticAnalysis.single().contains("src/App.tsx"))
    }

    @Test
    fun `non-writing review terminal receives the auto build gate`() {
        val planned = """
            {
              "goal": "implement then review",
              "stages": [
                { "id": "impl", "role": "implementer", "prompt": "write code", "produces": "code",
                  "kind": "react_entry", "writes": ["src/main.tsx"] },
                { "id": "review", "role": "reviewer", "prompt": "review against dod", "produces": "verdict",
                  "kind": "react_component", "needs": ["code"] }
              ],
              "edges": [
                { "from": "impl", "to": "review", "condition": { "type": "always_true" } },
                { "from": "review", "to": "done", "condition": { "type": "always_true" } }
              ]
            }
        """.trimIndent()
        val graph = codeCompiler.compile(planned, "wf")
        assertTrue(graph.stages.getValue(StageId("review")).autoBuildGate)
        // No dod in session → nothing to thread (execution plan never re-produces it).
        assertTrue(graph.stages.getValue(StageId("review")).needs.none { it.value == "dod" })

        // Planning-phase dod present in the session → threaded into impl + review needs (#264 seam).
        val withDod = codeCompiler.compile(planned, "wf", sessionArtifactIds = setOf("dod"))
        assertTrue(withDod.stages.getValue(StageId("review")).needs.any { it.value == "dod" })
        assertTrue(withDod.stages.getValue(StageId("impl")).needs.any { it.value == "dod" })
    }

    @Test
    fun `implementation and review consume dod only when the plan produces it`() {
        val planned = """
            {
              "goal": "define, implement, and review",
              "stages": [
                { "id": "analyst", "role": "analyst", "prompt": "define dod", "produces": "dod",
                  "kind": "react_component" },
                { "id": "impl", "role": "implementer", "prompt": "write code", "produces": "code",
                  "kind": "react_entry", "writes": ["src/main.tsx"] },
                { "id": "review", "role": "reviewer", "prompt": "review", "produces": "verdict",
                  "kind": "react_component", "needs": ["code"] }
              ],
              "edges": [
                { "from": "analyst", "to": "impl", "condition": { "type": "always_true" } },
                { "from": "impl", "to": "review", "condition": { "type": "always_true" } },
                { "from": "review", "to": "done", "condition": { "type": "always_true" } }
              ]
            }
        """.trimIndent()

        val graph = codeCompiler.compile(planned, "wf-with-dod")

        assertTrue(graph.stages.getValue(StageId("impl")).needs.any { it.value == "dod" })
        assertTrue(graph.stages.getValue(StageId("review")).needs.any { it.value == "dod" })
    }

    @Test
    fun `a MODULE declaration does not suppress the auto build gate`() {
        // MODULE is only a per-file typecheck (delegated to LSP, which can silently skip when
        // tsserver is missing), so it does NOT prove the assembled project builds — the terminal
        // whole-project floor must still be applied.
        val planned = """
            {
              "goal": "scaffold a react app",
              "stages": [
                { "id": "entry", "prompt": "write main.tsx", "produces": "app_entry", "kind": "react_entry",
                  "needs": [], "tools": ["file_write"], "writes": ["src/main.tsx"], "build_expectation": "module" },
                { "id": "views", "prompt": "write components", "produces": "app_views", "kind": "react_component",
                  "needs": ["app_entry"], "tools": ["file_write"], "writes": ["src/App.tsx"] }
              ],
              "edges": [
                { "from": "entry", "to": "views", "condition": { "type": "always_true" } },
                { "from": "views", "to": "done", "condition": { "type": "always_true" } }
              ]
            }
        """.trimIndent()
        val graph = codeCompiler.compile(planned, "wf")
        assertEquals(
            com.correx.core.transitions.graph.BuildExpectation.MODULE,
            graph.stages[StageId("entry")]!!.buildExpectation,
        )
        assertTrue(
            graph.stages.values.any { it.autoBuildGate },
            "MODULE (typecheck-only) must not remove the terminal project build floor",
        )
    }

    @Test
    fun `an explicitly declared PROJECT gate suppresses per-stage gates but not the terminal floor`() {
        // A mid-plan PROJECT build verifies only what existed when it ran, so it drops the
        // redundant per-writing-stage gates while the terminal floor stays (#263).
        val planned = """
            {
              "goal": "scaffold a react app",
              "stages": [
                { "id": "entry", "prompt": "write main.tsx", "produces": "app_entry", "kind": "react_entry",
                  "needs": [], "tools": ["file_write"], "writes": ["src/main.tsx"], "build_expectation": "project" },
                { "id": "views", "prompt": "write components", "produces": "app_views", "kind": "react_component",
                  "needs": ["app_entry"], "tools": ["file_write"], "writes": ["src/App.tsx"] }
              ],
              "edges": [
                { "from": "entry", "to": "views", "condition": { "type": "always_true" } },
                { "from": "views", "to": "done", "condition": { "type": "always_true" } }
              ]
            }
        """.trimIndent()
        val graph = codeCompiler.compile(planned, "wf")
        assertEquals(
            com.correx.core.transitions.graph.BuildExpectation.PROJECT,
            graph.stages[StageId("entry")]!!.buildExpectation,
        )
        assertFalse(
            graph.stages.getValue(StageId("entry")).autoBuildGate,
            "a declared PROJECT build suppresses the redundant per-writing-stage gates",
        )
        assertTrue(
            graph.stages.getValue(StageId("views")).autoBuildGate,
            "the terminal stage keeps its build floor — the declared build ran before later writes",
        )
    }

    @Test
    fun `plan without declared writes has no auto build gate`() {
        // A plan with no write declarations cannot promise compiler coverage. Runtime still
        // inspects actual files for an opted-in gate, but the compiler must not pretend this
        // documentation-shaped plan has implementation-stage verification.
        val graph = compiler.compile(validPlan, "wf")
        assertEquals(
            0,
            graph.stages.values.count { it.autoBuildGate },
            "only write-declaring stages receive an auto build gate",
        )
    }
}
