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.ExecutionPlanLockedEvent
import com.correx.core.events.events.ExecutionPlanRejectedEvent
import com.correx.core.events.types.SessionId
import com.correx.core.kernel.execution.WorkflowResult
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 org.junit.jupiter.api.Assertions.assertEquals
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" }"""

    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 `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")
    }
}
