package com.correx.infrastructure.workflow

import com.correx.core.artifacts.kind.ConfigArtifactKind
import com.correx.core.artifacts.kind.JsonSchema
import com.correx.core.artifacts.kind.TypedArtifactSlot
import com.correx.core.events.types.ArtifactId
import com.correx.core.events.types.StageId
import com.correx.core.events.types.TransitionId
import com.correx.core.tools.contract.ToolCapability
import com.correx.core.transitions.conditions.AlwaysTrue
import com.correx.core.transitions.graph.StageConfig
import com.correx.core.transitions.graph.TransitionEdge
import com.correx.core.transitions.graph.WorkflowGraph
import kotlin.test.Test
import kotlin.test.assertEquals
import kotlin.test.assertTrue

class CapabilityGapDetectorTest {

    private fun kind(id: String) =
        ConfigArtifactKind(id, JsonSchema(type = "object", properties = emptyMap(), additionalProperties = true))

    private fun stage(
        prompt: String,
        tools: Set<String> = emptySet(),
        produces: List<String> = emptyList(),
    ) = StageConfig(
        produces = produces.map { TypedArtifactSlot(ArtifactId(it), kind(it)) },
        allowedTools = tools,
        metadata = mapOf("promptInline" to prompt),
    )

    private fun graph(stages: Map<String, StageConfig>, start: String = stages.keys.first()) = WorkflowGraph(
        id = "test",
        stages = stages.mapKeys { StageId(it.key) },
        transitions = setOf(TransitionEdge(TransitionId("e0"), StageId(start), StageId("done"), AlwaysTrue)),
        start = StageId(start),
    )

    private val fileWriteTool = "file_write" to setOf(ToolCapability.FILE_WRITE)
    private val fileReadTool = "file_read" to setOf(ToolCapability.FILE_READ)
    private val shellTool = "shell" to setOf(ToolCapability.SHELL_EXEC)

    @Test
    fun `a stage that writes a file without FILE_WRITE is a gap`() {
        val g = graph(mapOf("impl" to stage(prompt = "Create the new config file", tools = emptySet())))
        val gaps = CapabilityGapDetector.detect(g, emptyMap())
        val gap = gaps.single()
        assertEquals("impl", gap.stageId)
        assertEquals(ToolCapability.FILE_WRITE, gap.impliedCapability)
        assertEquals("create", gap.phrase)
        assertTrue(gap.grantedCapabilities.isEmpty())
    }

    @Test
    fun `a stage granted FILE_WRITE has no gap for a write-implying brief`() {
        val g = graph(mapOf("impl" to stage(prompt = "Create the new config file", tools = setOf(fileWriteTool.first))))
        val gaps = CapabilityGapDetector.detect(g, mapOf(fileWriteTool))
        assertTrue(gaps.isEmpty(), "expected no gap, got $gaps")
    }

    @Test
    fun `the motivating case - a recovery style stage with read-only tools implies write`() {
        val prompt = "Read the gate output and edit the files to fix the reported cause, then save your work."
        val g = graph(mapOf("recovery" to stage(prompt = prompt, tools = setOf(fileReadTool.first))))
        val gaps = CapabilityGapDetector.detect(g, mapOf(fileReadTool))
        val writeGap = gaps.single { it.impliedCapability == ToolCapability.FILE_WRITE }
        assertEquals("save", writeGap.phrase)
        assertTrue(gaps.none { it.impliedCapability == ToolCapability.FILE_READ }, "read is granted, no gap expected")
    }

    @Test
    fun `run a build without SHELL_EXEC is a gap`() {
        val g = graph(mapOf("verify" to stage(prompt = "Run the build and report the result")))
        val gaps = CapabilityGapDetector.detect(g, emptyMap())
        assertTrue(gaps.any { it.impliedCapability == ToolCapability.SHELL_EXEC && it.phrase == "run" })
        assertTrue(gaps.any { it.impliedCapability == ToolCapability.SHELL_EXEC && it.phrase == "build" })
    }

    @Test
    fun `a stage with an unrelated brief and no tools has no gaps`() {
        val g = graph(mapOf("analyse" to stage(prompt = "Summarize the requirements at a high level")))
        val gaps = CapabilityGapDetector.detect(g, emptyMap())
        assertTrue(gaps.isEmpty(), "expected no gap, got $gaps")
    }

    @Test
    fun `a stage fully equipped for its brief has no gaps`() {
        val prompt = "Read the spec, run the build, and create the output file"
        val g = graph(
            mapOf(
                "full" to stage(
                    prompt = prompt,
                    tools = setOf(fileReadTool.first, shellTool.first, fileWriteTool.first),
                ),
            ),
        )
        val gaps = CapabilityGapDetector.detect(g, mapOf(fileReadTool, shellTool, fileWriteTool))
        assertTrue(gaps.isEmpty(), "expected no gap, got $gaps")
    }
}
