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 org.junit.jupiter.api.Assumptions.assumeTrue
import org.junit.jupiter.api.Test
import java.nio.file.Files
import java.nio.file.Path
import kotlin.io.path.exists
import kotlin.io.path.writeText
import kotlin.test.assertEquals
import kotlin.test.assertTrue

class ReviewLoopWorkflowTest {

    private val registry = DefaultArtifactKindRegistry().apply {
        register(
            ConfigArtifactKind(
                id = "review_report",
                schema = JsonSchema(
                    type = "object",
                    properties = mapOf("verdict" to JsonSchemaProperty(type = "string")),
                    required = listOf("verdict"),
                    additionalProperties = true,
                ),
                llmEmitted = true,
            ),
        )
    }
    private val loader = TomlWorkflowLoader(registry)

    private val toml = """
        id = "review_loop"
        start = "implement"

        [[stages]]
        id = "implement"
        prompt = "prompts/implement.md"
        produces = [{ name = "patch", kind = "file_written" }]
        allowed_tools = ["file_write"]
        max_retries = 3

        [[stages]]
        id = "review"
        prompt = "prompts/review.md"
        needs = ["patch"]
        produces = [{ name = "review_report", kind = "review_report" }]
        max_retries = 3

        [[transitions]]
        id = "implement-to-review"
        from = "implement"
        to = "review"
        condition_type = "artifact_validated"
        condition_artifact_id = "patch"

        [[transitions]]
        id = "review-approved"
        from = "review"
        to = "done"
        condition_type = "artifact_field_equals"
        condition_artifact_id = "review_report"
        condition_field = "verdict"
        condition_value = "approved"
        condition_operator = "eq"

        [[transitions]]
        id = "review-changes-requested"
        from = "review"
        to = "implement"
        condition_type = "artifact_field_equals"
        condition_artifact_id = "review_report"
        condition_field = "verdict"
        condition_value = "approved"
        condition_operator = "neq"
    """.trimIndent()

    @Test
    fun `review loop parses with verdict-gated edges`() {
        val path = Files.createTempFile("review_loop", ".toml").also { it.writeText(toml) }
        val graph = loader.load(path)

        val review = graph.stages[StageId("review")]!!
        assertTrue(review.produces.any { it.name.value == "review_report" && it.kind.llmEmitted })

        val approved = graph.transitions.first { it.to == StageId("done") }.condition as ArtifactFieldEquals
        assertEquals("verdict", approved.field)
        assertEquals("approved", approved.value)
        assertEquals(FieldOperator.EQ, approved.operator)

        val changes = graph.transitions.first { it.to == StageId("implement") }.condition as ArtifactFieldEquals
        assertEquals(FieldOperator.NEQ, changes.operator)
    }

    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 `shipped review_loop grants the task tools to both stages`() {
        val path = repoFile("examples/workflows/review_loop.toml")
        assumeTrue(path != null, "review_loop.toml not found from ${System.getProperty("user.dir")}")
        val graph = loader.load(path!!)

        // implement claims/submits and may create a task, alongside its file write.
        assertEquals(
            setOf(
                "file_write", "task_create", "task_update", "task_context", "task_search",
                "mcp__codebase-memory__search_code",
                "mcp__codebase-memory__get_code_snippet",
                "mcp__codebase-memory__trace_path",
            ),
            graph.stages[StageId("implement")]!!.allowedTools,
        )
        // review 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",
            ),
            graph.stages[StageId("review")]!!.allowedTools,
        )
    }
}
