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.ArtifactValidated
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 FreestylePlanningWorkflowTest {

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

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

    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
    }

    @Test
    fun `freestyle_planning toml parses with discovery start plus analyst and architect stages`() {
        val path = repoFile("examples/workflows/freestyle_planning.toml")
        assumeTrue(path != null, "freestyle_planning.toml not found from ${System.getProperty("user.dir")}")

        val graph = TomlWorkflowLoader(registry).load(path!!)

        assertEquals(StageId("discovery"), graph.start)
        assertEquals(
            setOf("discovery", "analyst", "architect"),
            graph.stages.keys.map { it.value }.toSet(),
        )

        val discoveryToAnalyst = graph.transitions.first { it.from == StageId("discovery") }
        assertEquals(StageId("analyst"), discoveryToAnalyst.to)
        assertTrue(discoveryToAnalyst.condition is ArtifactValidated)

        val architect = graph.stages[StageId("architect")]!!
        assertTrue(
            architect.produces.any { it.name.value == "execution_plan" && it.kind.llmEmitted },
            "architect must produce execution_plan",
        )

        assertTrue(architect.needs.map { it.value }.contains("dod"), "architect needs the fixed DoD")

        val analystToArchitect = graph.transitions.first { it.from == StageId("analyst") }
        assertEquals(StageId("architect"), analystToArchitect.to)
        assertTrue(analystToArchitect.condition is ArtifactValidated)

        val architectToDone = graph.transitions.first { it.from == StageId("architect") }
        assertEquals(StageId("done"), architectToDone.to)
        assertTrue(architectToDone.condition is ArtifactValidated)

        assertEquals(3, graph.transitions.size)

        // analyst frames the work: search + open one task or decompose into a graph (the architect
        // threads the named task into the plan).
        assertEquals(
            setOf("file_read", "list_dir", "shell", "task_search", "task_context", "task_create", "task_decompose"),
            graph.stages[StageId("analyst")]!!.allowedTools,
        )
    }
}
