package com.correx.apps.server.freestyle

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.events.events.CapabilityGapDetectedEvent
import com.correx.core.events.events.CapabilityGapReflectedEvent
import com.correx.core.events.events.CapabilityGapVerdict
import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.ExecutionPlanLockedEvent
import com.correx.core.events.events.ExecutionPlanRejectedEvent
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.PlanLintCompletedEvent
import com.correx.core.events.events.RepoMapComputedEvent
import com.correx.core.events.events.RepoMapEntry
import com.correx.core.events.types.EventId
import com.correx.core.events.types.SessionId
import com.correx.core.kernel.execution.WorkflowResult
import com.correx.core.kernel.orchestration.CapabilityGapReflection
import com.correx.core.kernel.orchestration.CapabilityGapReflector
import com.correx.core.kernel.orchestration.OrchestrationConfig
import com.correx.core.transitions.graph.WorkflowGraph
import com.correx.infrastructure.persistence.InMemoryEventStore
import com.correx.infrastructure.workflow.ExecutionPlanCompiler
import kotlinx.coroutines.runBlocking
import kotlinx.datetime.Clock
import java.util.UUID
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test

class FreestyleDriverTest {

    private val validPlanJson = """
        {
          "goal": "test plan",
          "stages": [
            {
              "id": "analyse",
              "prompt": "Analyse the problem",
              "produces": "patch",
              "needs": [],
              "tools": []
            },
            {
              "id": "apply",
              "prompt": "Apply the fix",
              "produces": "patch",
              "needs": ["patch"],
              "tools": []
            }
          ],
          "edges": [
            {
              "from": "analyse",
              "to": "apply",
              "condition": { "type": "always_true" }
            },
            {
              "from": "apply",
              "to": "done",
              "condition": { "type": "always_true" }
            }
          ]
        }
    """.trimIndent()

    private val malformedPlanJson = """{ "not": "a valid plan" }"""

    // Compiles cleanly, but "apply" needs "ghost", which no stage produces — a lint hard failure
    // (the compiler does not check needs satisfaction).
    private val lintFailingPlanJson = """
        {
          "goal": "broken plan",
          "stages": [
            { "id": "analyse", "prompt": "Analyse", "produces": "patch", "needs": [], "tools": [] },
            { "id": "apply", "prompt": "Apply", "produces": "patch", "needs": ["ghost"], "tools": [] }
          ],
          "edges": [
            { "from": "analyse", "to": "apply", "condition": { "type": "always_true" } },
            { "from": "apply", "to": "done", "condition": { "type": "always_true" } }
          ]
        }
    """.trimIndent()

    private fun buildRegistry(): DefaultArtifactKindRegistry =
        DefaultArtifactKindRegistry().also {
            it.register(ConfigArtifactKind(id = "patch", schema = JsonSchema(type = "object"), llmEmitted = true))
        }

    @Test
    fun `lockAndRun emits ExecutionPlanLockedEvent with compiled stageIds and invokes runPhase2`(): Unit = runBlocking {
        val sessionId = SessionId("driver-test-session")
        val eventStore = InMemoryEventStore()
        val registry = buildRegistry()
        val compiler = ExecutionPlanCompiler(registry)

        var runPhase2Invocations = 0
        var capturedGraph: WorkflowGraph? = null

        val driver = FreestyleDriver(
            eventStore = eventStore,
            compiler = compiler,
            planContent = { validPlanJson },
            config = OrchestrationConfig(),
            runPhase2 = { sid, graph, _ ->
                runPhase2Invocations++
                capturedGraph = graph
                WorkflowResult.Completed(sid, graph.start)
            },
        )

        driver.lockAndRun(sessionId)

        val lockedEvents = eventStore.read(sessionId)
            .map { it.payload }
            .filterIsInstance<ExecutionPlanLockedEvent>()

        assertEquals(1, lockedEvents.size, "Expected exactly one ExecutionPlanLockedEvent")
        val locked = lockedEvents.single()
        assertEquals(setOf("analyse", "apply"), locked.stageIds.toSet())
        assertEquals("analyse", locked.startStageId)
        assertTrue(locked.workflowId.startsWith("freestyle-"))
        assertEquals(1, runPhase2Invocations, "runPhase2 should be called exactly once")
        assertEquals(setOf("analyse", "apply"), capturedGraph?.stageIds?.map { it.value }?.toSet())
    }

    @Test
    fun `lockAndRun with malformed plan emits no lock event and does not invoke runPhase2`(): Unit = runBlocking {
        val sessionId = SessionId("driver-malformed-session")
        val eventStore = InMemoryEventStore()
        val registry = buildRegistry()
        val compiler = ExecutionPlanCompiler(registry)

        var runPhase2Invocations = 0

        val driver = FreestyleDriver(
            eventStore = eventStore,
            compiler = compiler,
            planContent = { malformedPlanJson },
            config = OrchestrationConfig(),
            runPhase2 = { _, _, _ ->
                runPhase2Invocations++
                WorkflowResult.Completed(sessionId, com.correx.core.events.types.StageId("done"))
            },
        )

        driver.lockAndRun(sessionId)

        val lockedEvents = eventStore.read(sessionId)
            .map { it.payload }
            .filterIsInstance<ExecutionPlanLockedEvent>()

        assertTrue(lockedEvents.isEmpty(), "No lock event expected for malformed plan")
        assertEquals(0, runPhase2Invocations, "runPhase2 must not be called for malformed plan")

        val rejectedEvents = eventStore.read(sessionId)
            .map { it.payload }
            .filterIsInstance<ExecutionPlanRejectedEvent>()
        assertEquals(1, rejectedEvents.size, "Expected exactly one ExecutionPlanRejectedEvent")
        assertEquals("compile", rejectedEvents.single().source)
    }

    @Test
    fun `lockAndRun with absent plan content does not emit lock event or invoke runPhase2`(): Unit = runBlocking {
        val sessionId = SessionId("driver-absent-session")
        val eventStore = InMemoryEventStore()
        val registry = buildRegistry()
        val compiler = ExecutionPlanCompiler(registry)

        var runPhase2Invocations = 0

        val driver = FreestyleDriver(
            eventStore = eventStore,
            compiler = compiler,
            planContent = { null },
            config = OrchestrationConfig(),
            runPhase2 = { _, _, _ ->
                runPhase2Invocations++
                WorkflowResult.Completed(sessionId, com.correx.core.events.types.StageId("done"))
            },
        )

        driver.lockAndRun(sessionId)

        val lockedEvents = eventStore.read(sessionId)
            .map { it.payload }
            .filterIsInstance<ExecutionPlanLockedEvent>()

        assertTrue(lockedEvents.isEmpty(), "No lock event when plan content is absent")
        assertEquals(0, runPhase2Invocations, "runPhase2 must not be called when plan is absent")

        val rejectedEvents = eventStore.read(sessionId)
            .map { it.payload }
            .filterIsInstance<ExecutionPlanRejectedEvent>()
        assertEquals(1, rejectedEvents.size, "Expected exactly one ExecutionPlanRejectedEvent")
        assertEquals("missing_content", rejectedEvents.single().source)
    }

    @Test
    fun `operator rejection at plan gate emits rejected and skips lock and phase2`(): Unit = runBlocking {
        val sessionId = SessionId("driver-operator-reject-session")
        val eventStore = InMemoryEventStore()
        val registry = buildRegistry()
        val compiler = ExecutionPlanCompiler(registry)

        var runPhase2Invocations = 0

        val driver = FreestyleDriver(
            eventStore = eventStore,
            compiler = compiler,
            planContent = { validPlanJson },
            config = OrchestrationConfig(),
            runPhase2 = { _, _, _ ->
                runPhase2Invocations++
                WorkflowResult.Completed(sessionId, com.correx.core.events.types.StageId("done"))
            },
            requestPlanApproval = { _, _ -> false },
        )

        driver.lockAndRun(sessionId)

        val lockedEvents = eventStore.read(sessionId)
            .map { it.payload }
            .filterIsInstance<ExecutionPlanLockedEvent>()

        assertTrue(lockedEvents.isEmpty(), "No lock event expected when operator rejects")
        assertEquals(0, runPhase2Invocations, "runPhase2 must not be called when operator rejects")

        val rejectedEvents = eventStore.read(sessionId)
            .map { it.payload }
            .filterIsInstance<ExecutionPlanRejectedEvent>()
        assertEquals(1, rejectedEvents.size, "Expected exactly one ExecutionPlanRejectedEvent")
        assertEquals("operator", rejectedEvents.single().source)
    }

    @Test
    fun `plan failing lint is rejected with source lint and never locks or runs phase2`(): Unit = runBlocking {
        val sessionId = SessionId("driver-lint-fail-session")
        val eventStore = InMemoryEventStore()
        val compiler = ExecutionPlanCompiler(buildRegistry())

        var runPhase2Invocations = 0
        var approvalRequested = false

        val driver = FreestyleDriver(
            eventStore = eventStore,
            compiler = compiler,
            planContent = { lintFailingPlanJson },
            config = OrchestrationConfig(),
            runPhase2 = { _, _, _ ->
                runPhase2Invocations++
                WorkflowResult.Completed(sessionId, com.correx.core.events.types.StageId("done"))
            },
            requestPlanApproval = { _, _ -> approvalRequested = true; true },
        )

        driver.lockAndRun(sessionId)

        val payloads = eventStore.read(sessionId).map { it.payload }
        assertTrue(payloads.filterIsInstance<ExecutionPlanLockedEvent>().isEmpty(), "broken plan must not lock")
        assertEquals(0, runPhase2Invocations, "runPhase2 must not run on a lint failure")
        assertFalse(approvalRequested, "operator must not be asked to approve a plan that already failed lint")

        val lint = payloads.filterIsInstance<PlanLintCompletedEvent>().single()
        assertTrue(
            lint.hardFailures.any { it.code == "unproduced_need" },
            "expected unproduced_need: ${lint.hardFailures}",
        )

        val rejected = payloads.filterIsInstance<ExecutionPlanRejectedEvent>().single()
        assertEquals("lint", rejected.source)
    }

    @Test
    fun `a clean plan records a passing PlanLintCompletedEvent before locking`(): Unit = runBlocking {
        val sessionId = SessionId("driver-lint-clean-session")
        val eventStore = InMemoryEventStore()
        val compiler = ExecutionPlanCompiler(buildRegistry())

        val driver = FreestyleDriver(
            eventStore = eventStore,
            compiler = compiler,
            planContent = { validPlanJson },
            config = OrchestrationConfig(),
            runPhase2 = { sid, graph, _ -> WorkflowResult.Completed(sid, graph.start) },
        )

        driver.lockAndRun(sessionId)

        val payloads = eventStore.read(sessionId).map { it.payload }
        val lint = payloads.filterIsInstance<PlanLintCompletedEvent>().single()
        assertTrue(lint.hardFailures.isEmpty(), "clean plan should have no hard failures: ${lint.hardFailures}")
        assertEquals(1, payloads.filterIsInstance<ExecutionPlanLockedEvent>().size, "clean plan should lock")
    }

    @Test
    fun `operator approval at plan gate locks and runs phase2`(): Unit = runBlocking {
        val sessionId = SessionId("driver-operator-approve-session")
        val eventStore = InMemoryEventStore()
        val registry = buildRegistry()
        val compiler = ExecutionPlanCompiler(registry)

        var runPhase2Invocations = 0
        var approvalPreview: String? = null

        val driver = FreestyleDriver(
            eventStore = eventStore,
            compiler = compiler,
            planContent = { validPlanJson },
            config = OrchestrationConfig(),
            runPhase2 = { sid, graph, _ ->
                runPhase2Invocations++
                WorkflowResult.Completed(sid, graph.start)
            },
            requestPlanApproval = { _, preview ->
                approvalPreview = preview
                assertTrue(preview.contains("stages"), "Preview should contain plan JSON with stages")
                true
            },
        )

        driver.lockAndRun(sessionId)

        val lockedEvents = eventStore.read(sessionId)
            .map { it.payload }
            .filterIsInstance<ExecutionPlanLockedEvent>()

        assertEquals(1, lockedEvents.size, "Expected exactly one ExecutionPlanLockedEvent after approval")
        assertEquals(1, runPhase2Invocations, "runPhase2 should be called exactly once after approval")
        assertTrue(approvalPreview != null, "requestPlanApproval should have been invoked")
    }

    // A stage brief that trips the FILE_WRITE gap (verb "create") with zero granted tools, so
    // CapabilityGapDetector.detect() always finds exactly one gap on stage "apply" for these tests.
    private val gapPlanJson = """
        {
          "goal": "test plan with a capability gap",
          "stages": [
            { "id": "analyse", "prompt": "Analyse the problem", "produces": "patch", "needs": [], "tools": [] },
            { "id": "apply", "prompt": "create the output file", "produces": "patch", "needs": ["patch"], "tools": [] }
          ],
          "edges": [
            { "from": "analyse", "to": "apply", "condition": { "type": "always_true" } },
            { "from": "apply", "to": "done", "condition": { "type": "always_true" } }
          ]
        }
    """.trimIndent()

    @Test
    fun `a RESOLVED reflection is advisory only and does not touch the approval preview`(): Unit = runBlocking {
        val sessionId = SessionId("driver-reflect-resolved-session")
        val eventStore = InMemoryEventStore()
        val compiler = ExecutionPlanCompiler(buildRegistry())
        var approvalPreview: String? = null

        val reflector = CapabilityGapReflector { _, gaps ->
            gaps.map {
                CapabilityGapReflection(
                    it.stageId, it.impliedCapability, CapabilityGapVerdict.RESOLVED, "revised: reuse file_write",
                )
            }
        }

        val driver = FreestyleDriver(
            eventStore = eventStore,
            compiler = compiler,
            planContent = { gapPlanJson },
            config = OrchestrationConfig(),
            runPhase2 = { sid, graph, _ -> WorkflowResult.Completed(sid, graph.start) },
            requestPlanApproval = { _, preview -> approvalPreview = preview; true },
            reflector = reflector,
        )

        driver.lockAndRun(sessionId)

        val payloads = eventStore.read(sessionId).map { it.payload }
        assertEquals(1, payloads.filterIsInstance<CapabilityGapDetectedEvent>().size)
        val reflected = payloads.filterIsInstance<CapabilityGapReflectedEvent>().single()
        assertEquals(CapabilityGapVerdict.RESOLVED, reflected.verdict)
        assertEquals(1, payloads.filterIsInstance<ExecutionPlanLockedEvent>().size, "RESOLVED must not block locking")
        assertFalse(
            approvalPreview.orEmpty().contains("CAPABILITY GAPS REQUESTED"),
            "RESOLVED verdicts must not be surfaced as an escalation",
        )
    }

    @Test
    fun `a NEEDS_TOOL reflection is surfaced to the approval seam but never auto-grants`(): Unit = runBlocking {
        val sessionId = SessionId("driver-reflect-needs-tool-session")
        val eventStore = InMemoryEventStore()
        val compiler = ExecutionPlanCompiler(buildRegistry())
        var approvalPreview: String? = null

        val reflector = CapabilityGapReflector { _, gaps ->
            gaps.map {
                CapabilityGapReflection(
                    it.stageId, it.impliedCapability, CapabilityGapVerdict.NEEDS_TOOL, "genuinely needs file_write",
                )
            }
        }

        val driver = FreestyleDriver(
            eventStore = eventStore,
            compiler = compiler,
            planContent = { gapPlanJson },
            config = OrchestrationConfig(),
            runPhase2 = { sid, graph, _ -> WorkflowResult.Completed(sid, graph.start) },
            requestPlanApproval = { _, preview -> approvalPreview = preview; true },
            reflector = reflector,
        )

        driver.lockAndRun(sessionId)

        val payloads = eventStore.read(sessionId).map { it.payload }
        val reflected = payloads.filterIsInstance<CapabilityGapReflectedEvent>().single()
        assertEquals(CapabilityGapVerdict.NEEDS_TOOL, reflected.verdict)
        assertTrue(
            approvalPreview.orEmpty().contains("CAPABILITY GAPS REQUESTED"),
            "NEEDS_TOOL must be surfaced into the human approval preview",
        )
        assertTrue(approvalPreview.orEmpty().contains("apply"), "escalation note should name the stage")
        // No tool grant mechanism exists on this seam at all — the plan still locks purely on the
        // operator's approval decision, exactly as an ungapped plan would.
        assertEquals(1, payloads.filterIsInstance<ExecutionPlanLockedEvent>().size)
    }

    @Test
    fun `malformed reflection output degrades to advisory and never crashes lockAndRun`(): Unit = runBlocking {
        val sessionId = SessionId("driver-reflect-malformed-session")
        val eventStore = InMemoryEventStore()
        val compiler = ExecutionPlanCompiler(buildRegistry())

        val reflector = CapabilityGapReflector { _, _ -> throw IllegalStateException("boom") }

        val driver = FreestyleDriver(
            eventStore = eventStore,
            compiler = compiler,
            planContent = { gapPlanJson },
            config = OrchestrationConfig(),
            runPhase2 = { sid, graph, _ -> WorkflowResult.Completed(sid, graph.start) },
            reflector = reflector,
        )

        driver.lockAndRun(sessionId)

        val payloads = eventStore.read(sessionId).map { it.payload }
        assertTrue(
            payloads.filterIsInstance<CapabilityGapReflectedEvent>().isEmpty(),
            "no verdict when the judge itself throws",
        )
        assertEquals(
            1,
            payloads.filterIsInstance<ExecutionPlanLockedEvent>().size,
            "a broken reflector must not wedge the plan gate",
        )
    }

    @Test
    fun `no reflector wired preserves part-1 behavior`(): Unit = runBlocking {
        val sessionId = SessionId("driver-no-reflector-session")
        val eventStore = InMemoryEventStore()
        val compiler = ExecutionPlanCompiler(buildRegistry())

        val driver = FreestyleDriver(
            eventStore = eventStore,
            compiler = compiler,
            planContent = { gapPlanJson },
            config = OrchestrationConfig(),
            runPhase2 = { sid, graph, _ -> WorkflowResult.Completed(sid, graph.start) },
        )

        driver.lockAndRun(sessionId)

        val payloads = eventStore.read(sessionId).map { it.payload }
        assertEquals(1, payloads.filterIsInstance<CapabilityGapDetectedEvent>().size)
        assertTrue(payloads.filterIsInstance<CapabilityGapReflectedEvent>().isEmpty())
        assertEquals(1, payloads.filterIsInstance<ExecutionPlanLockedEvent>().size)
    }

    // "apply" is scoped to frontend/** — with a scan recorded whose only path is backend/, and no
    // stage creating a frontend file, scope grounding fails (RETURN_TO_ARCHITECT). Passes lint.
    private val groundingFailPlanJson = """
        {
          "goal": "plan scoped to a path that doesn't exist",
          "stages": [
            { "id": "analyse", "prompt": "Analyse", "produces": "patch", "needs": [], "tools": [] },
            { "id": "apply", "prompt": "Apply", "produces": "patch", "needs": ["patch"], "tools": [],
              "touches": ["frontend/**"] }
          ],
          "edges": [
            { "from": "analyse", "to": "apply", "condition": { "type": "always_true" } },
            { "from": "apply", "to": "done", "condition": { "type": "always_true" } }
          ]
        }
    """.trimIndent()

    // Records a scan (scanned=true) whose only path is backend/, so frontend/** grounds to nothing.
    private suspend fun recordBackendScan(store: InMemoryEventStore, sessionId: SessionId) {
        store.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = Clock.System.now(),
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = RepoMapComputedEvent(
                    sessionId = sessionId,
                    repoRoot = "/ws",
                    entries = listOf(RepoMapEntry(path = "backend/main.kt", score = 1.0)),
                    computedAt = Clock.System.now(),
                ),
            ),
        )
    }

    @Test
    fun `grounding failure returns plan to architect then locks once the re-run yields a grounded plan`(): Unit =
        runBlocking {
            val sessionId = SessionId("driver-grounding-retry-session")
            val eventStore = InMemoryEventStore()
            val compiler = ExecutionPlanCompiler(buildRegistry())
            recordBackendScan(eventStore, sessionId)

            var corrected = false
            var rerunInvocations = 0
            var runPhase2Invocations = 0

            val driver = FreestyleDriver(
                eventStore = eventStore,
                compiler = compiler,
                // Fails grounding until the architect re-run "fixes" it (flips to the unconstrained plan).
                planContent = { if (corrected) validPlanJson else groundingFailPlanJson },
                config = OrchestrationConfig(),
                runPhase2 = { sid, graph, _ ->
                    runPhase2Invocations++
                    WorkflowResult.Completed(sid, graph.start)
                },
                rerunArchitect = { sid ->
                    rerunInvocations++
                    corrected = true
                    WorkflowResult.Completed(sid, com.correx.core.events.types.StageId("architect"))
                },
            )

            driver.lockAndRun(sessionId)

            assertEquals(1, rerunInvocations, "architect should be re-run exactly once")
            val payloads = eventStore.read(sessionId).map { it.payload }
            assertEquals(1, payloads.filterIsInstance<ExecutionPlanLockedEvent>().size, "corrected plan should lock")
            assertEquals(1, runPhase2Invocations, "runPhase2 should run once on the corrected plan")
            assertTrue(
                payloads.filterIsInstance<ExecutionPlanRejectedEvent>().isEmpty(),
                "no rejection once the re-run grounds",
            )
        }

    @Test
    fun `grounding failure that never resolves is rejected with source grounding after exhausting retries`(): Unit =
        runBlocking {
            val sessionId = SessionId("driver-grounding-exhaust-session")
            val eventStore = InMemoryEventStore()
            val compiler = ExecutionPlanCompiler(buildRegistry())
            recordBackendScan(eventStore, sessionId)

            var rerunInvocations = 0
            var runPhase2Invocations = 0

            val driver = FreestyleDriver(
                eventStore = eventStore,
                compiler = compiler,
                planContent = { groundingFailPlanJson }, // never corrected
                config = OrchestrationConfig(),
                runPhase2 = { sid, graph, _ ->
                    runPhase2Invocations++
                    WorkflowResult.Completed(sid, graph.start)
                },
                rerunArchitect = { sid ->
                    rerunInvocations++
                    WorkflowResult.Completed(sid, com.correx.core.events.types.StageId("architect"))
                },
                maxGroundingRetries = 2,
            )

            driver.lockAndRun(sessionId)

            assertEquals(2, rerunInvocations, "architect re-run should be capped at maxGroundingRetries")
            assertEquals(0, runPhase2Invocations, "runPhase2 must not run when grounding never clears")
            val payloads = eventStore.read(sessionId).map { it.payload }
            assertTrue(payloads.filterIsInstance<ExecutionPlanLockedEvent>().isEmpty(), "never locks")
            val rejected = payloads.filterIsInstance<ExecutionPlanRejectedEvent>().single()
            assertEquals("grounding", rejected.source)
        }
}
