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 com.correx.core.transitions.conditions.ArtifactFieldEquals
import com.correx.core.transitions.conditions.FieldOperator
import com.correx.core.transitions.conditions.TasksReady
import org.junit.jupiter.api.Assumptions.assumeTrue
import org.junit.jupiter.api.Test
import java.nio.file.Path
import kotlin.io.path.exists
import kotlin.test.assertEquals
import kotlin.test.assertTrue

class RolePipelineWorkflowTest {

    private fun llmKind(id: String) = ConfigArtifactKind(
        id = id,
        schema = JsonSchema(
            type = "object",
            properties = mapOf("verdict" to JsonSchemaProperty(type = "string")),
            additionalProperties = true,
        ),
        llmEmitted = true,
    )

    private val registry = DefaultArtifactKindRegistry().apply {
        register(llmKind("discovery"))
        register(llmKind("analysis"))
        register(llmKind("design"))
        register(llmKind("review_report"))
    }

    // Walk up from the working dir to the repo root so the test finds the shipped file
    // regardless of which module dir Gradle runs it from.
    private fun repoFile(rel: String): Path? {
        var dir: Path? = Path.of(System.getProperty("user.dir")).toAbsolutePath()
        while (dir != null) {
            val candidate = dir.resolve(rel)
            if (candidate.exists()) return candidate
            dir = dir.parent
        }
        return null
    }

    private fun loadGraph() =
        repoFile("examples/workflows/role_pipeline.toml")?.let { TomlWorkflowLoader(registry).load(it) }

    @Test
    fun `shipped role_pipeline toml parses with the tasked execution graph`() {
        val graph = loadGraph()
        assumeTrue(graph != null, "role_pipeline.toml not found from ${System.getProperty("user.dir")}")
        graph!!

        assertEquals(StageId("discovery"), graph.start)
        assertEquals(
            setOf("discovery", "analyst", "architect", "decomposer", "implementer", "reviewer"),
            graph.stages.keys.map { it.value }.toSet(),
        )
        assertEquals(8, graph.transitions.size)

        // discovery is the read-only clarification gate that precedes the analyst.
        val discoveryToAnalyst = graph.transitions.first { it.from == StageId("discovery") }
        assertEquals(StageId("analyst"), discoveryToAnalyst.to)

        // The decomposer is the mandatory task-graph gate.
        assertEquals("true", graph.stages[StageId("decomposer")]!!.metadata["requireTaskDecompose"])
        // The implementer claims a task on entry.
        assertEquals("true", graph.stages[StageId("implementer")]!!.metadata["claimTask"])
    }

    @Test
    fun `loop is gated by verdict and the task board`() {
        val graph = loadGraph()
        assumeTrue(graph != null, "role_pipeline.toml not found")
        graph!!

        // Decomposer only advances once tasks exist.
        val toImpl = graph.transitions
            .first { it.from == StageId("decomposer") && it.to == StageId("implementer") }
        assertEquals(TasksReady, toImpl.condition)

        // Reviewer edges, in id (= evaluation) order: changes → impl, tasks_ready → impl, else → done.
        val reviewerEdges = graph.transitions
            .filter { it.from == StageId("reviewer") }
            .sortedBy { it.id.value }
        assertEquals(
            listOf("review-1-changes", "review-2-more", "review-3-done"),
            reviewerEdges.map { it.id.value },
        )
        // 1. changes requested (verdict != approved) loops back to the implementer.
        val changes = reviewerEdges[0]
        assertEquals(StageId("implementer"), changes.to)
        assertEquals(FieldOperator.NEQ, (changes.condition as ArtifactFieldEquals).operator)
        // 2. approved + tasks remain → claim the next.
        assertEquals(StageId("implementer"), reviewerEdges[1].to)
        assertEquals(TasksReady, reviewerEdges[1].condition)
        // 3. approved + board drained → done (fallthrough).
        assertEquals(StageId("done"), reviewerEdges[2].to)
    }

    @Test
    fun `task tools are granted to the right stages`() {
        val graph = loadGraph()
        assumeTrue(graph != null, "role_pipeline.toml not found")
        graph!!

        // implementer owns the lifecycle: claim/submit + the full file set + the scope valve.
        val implTools = graph.stages[StageId("implementer")]!!.allowedTools
        assertTrue("propose_scope" in implTools, "implementer must be able to propose scope widening")
        assertTrue(implTools.containsAll(setOf("file_write", "file_edit", "task_update")))
        // reviewer reads the task and completes it on an approved verdict.
        assertEquals(
            setOf(
                "task_context", "task_update",
                "mcp__codebase-memory__search_code",
                "mcp__codebase-memory__trace_path",
                "mcp__codebase-memory__get_code_snippet",
            ),
            graph.stages[StageId("reviewer")]!!.allowedTools,
        )
    }
}
