package com.correx.apps.server.freestyle

import com.correx.core.events.events.ArtifactValidatedEvent
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.PlanGroundingEvaluatedEvent
import com.correx.core.events.events.PlanGroundingVerdict
import com.correx.core.events.events.PlanLintCompletedEvent
import com.correx.core.events.events.ProjectProfileBoundEvent
import com.correx.core.events.events.RepoMapComputedEvent
import com.correx.core.events.events.WorkflowFailedEvent
import com.correx.infrastructure.workflow.PlanGrounder
import com.correx.core.tools.contract.ToolCapability
import com.correx.infrastructure.workflow.CapabilityGap
import com.correx.infrastructure.workflow.CapabilityGapDetector
import com.correx.infrastructure.workflow.PlanLintResult
import com.correx.infrastructure.workflow.PlanLinter
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ArtifactId
import com.correx.core.events.types.EventId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.kernel.execution.WorkflowResult
import com.correx.core.kernel.orchestration.CapabilityGapInput
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.workflow.ExecutionPlanCompiler
import kotlinx.datetime.Clock
import org.slf4j.LoggerFactory
import java.util.UUID

/**
 * Drives the second phase of a freestyle workflow: reads the validated execution_plan,
 * compiles it, emits [ExecutionPlanLockedEvent], and invokes phase-2 orchestration.
 *
 * [runPhase2] is a functional seam so the driver can be tested without the full
 * [DefaultSessionOrchestrator] constructor chain.
 */
@Suppress("LongParameterList")
class FreestyleDriver(
    private val eventStore: EventStore,
    private val compiler: ExecutionPlanCompiler,
    private val planContent: (SessionId) -> String?,
    private val config: OrchestrationConfig,
    private val runPhase2: suspend (SessionId, WorkflowGraph, OrchestrationConfig) -> WorkflowResult,
    private val requestPlanApproval: suspend (SessionId, String) -> Boolean = { _, _ -> true },
    // Tool name -> the capabilities that tool grants (Tool.requiredCapabilities), used only by the
    // capability-gap detector below. Empty = detector finds nothing (preserves callers/tests that
    // don't care about it) rather than depending on a live ToolRegistry here.
    private val toolCapabilities: Map<String, Set<ToolCapability>> = emptyMap(),
    // Vikunja #30 part 2: bounded LLM "are you sure?" pass over capability gaps, consulted before
    // requestPlanApproval. Null = feature degrades to part-1 behavior (gaps recorded, never reflected).
    private val reflector: CapabilityGapReflector? = null,
    // Return-to-architect loop: on a grounding rejection, re-run the planning workflow from the
    // architect stage so it emits a corrected plan (the grounding findings are injected into its
    // context by buildGroundingFeedbackEntry). Null = no loop; grounding rejection is terminal.
    private val rerunArchitect: (suspend (SessionId) -> WorkflowResult)? = null,
    // Max grounding-driven architect re-runs before the plan is rejected for good.
    private val maxGroundingRetries: Int = 2,
) {
    @Suppress("ReturnCount") // sequential gate pipeline: each gate is a guard-return, same as the engines
    suspend fun lockAndRun(sessionId: SessionId) {
        // Return-to-architect loop: each pass compiles + gates the current execution_plan. A grounding
        // rejection with budget left re-runs the architect (which emits a corrected plan) and loops;
        // every other gate failure — and grounding once the budget is spent — is terminal.
        var groundingRetries = 0
        while (true) {
            val json = planContent(sessionId) ?: run {
                log.warn("freestyle: no execution_plan content for session={}", sessionId.value)
                emitRejected(sessionId, "no execution_plan content produced by planning phase", "missing_content")
                return
            }
            val graph = runCatching {
                compiler.compile(json, "freestyle-${sessionId.value}", sessionArtifacts(sessionId))
            }
                .getOrElse {
                    log.error("freestyle: plan failed to compile: {}", it.message)
                    emitRejected(sessionId, "plan failed to compile: ${it.message}", "compile")
                    return
                }
            // Deterministic lint (plan-pipeline-spec §5) before the plan is surfaced/locked: a hard
            // failure (unproduced need, trap state) means the plan would fail at runtime, so reject it
            // now rather than execute a broken plan. Soft findings are recorded for display only.
            val lint = PlanLinter.lint(graph)
            emitPlanLint(sessionId, graph.id, lint)
            if (!lint.passed) {
                val summary = lint.hardFailures.joinToString("; ") {
                    "${it.code}${it.stageId?.let { s -> " @$s" } ?: ""}: ${it.detail}"
                }
                log.warn("freestyle: plan failed lint for session={}: {}", sessionId.value, summary)
                emitRejected(sessionId, "plan failed lint: $summary", "lint")
                return
            }
            // Plan grounding (design 2026-07-15 seam 1) before lock: check the compiled plan against the
            // session's recorded workspace facts (repo map + profile commands). A build stage aimed at a
            // prerequisite nothing creates is doomed. Deterministic + pure over recorded events (#8/#9).
            val groundingFailure = groundPlan(sessionId, graph)
            if (groundingFailure != null) {
                if (groundingRetries < maxGroundingRetries && rerunArchitect != null) {
                    groundingRetries++
                    log.info(
                        "freestyle: grounding returned plan to architect (attempt {}/{}) session={}: {}",
                        groundingRetries, maxGroundingRetries, sessionId.value, groundingFailure,
                    )
                    // Re-run the architect stage; it sees the recorded grounding findings via
                    // buildGroundingFeedbackEntry and emits a corrected plan. Then loop and re-gate.
                    rerunArchitect.invoke(sessionId)
                    continue
                }
                log.warn("freestyle: plan failed grounding for session={}: {}", sessionId.value, groundingFailure)
                emitRejected(sessionId, "plan failed grounding: $groundingFailure", "grounding")
                return
            }
            lockAndRunGrounded(sessionId, graph, json)
            return
        }
    }

    /** Post-grounding tail of [lockAndRun]: capability gaps, plan-approval, lock, phase-2 handoff. */
    private suspend fun lockAndRunGrounded(sessionId: SessionId, graph: WorkflowGraph, json: String) {
        // Capability-gap detector (Vikunja #30 part 1): advisory only — recorded, never blocking.
        // A gap does not fail the gate and does not grant the missing tool (invariants #3/#4/#5).
        val gaps = CapabilityGapDetector.detect(graph, toolCapabilities)
        gaps.forEach { gap ->
            log.info(
                "[FreestyleDriver] capability gap session={} stage={} implied={} evidence='{}'",
                sessionId.value, gap.stageId, gap.impliedCapability, gap.phrase,
            )
            emitCapabilityGap(sessionId, gap.stageId, gap.impliedCapability.name, gap.phrase)
        }
        // Reflection rung (Vikunja #30 part 2): one bounded LLM pass over the whole batch of gaps
        // before the human is bothered. RESOLVED is advisory-only (logged, never rewrites the locked
        // plan — invariant #3 LLM-proposes/core-decides); NEEDS_TOOL is the only case surfaced into
        // the plan-review approval preview below, since that is the existing human approval seam and
        // a tool grant must never happen without it (invariants #4/#5).
        val needsToolNote = reflectOnGaps(sessionId, gaps)
        val approved = requestPlanApproval(sessionId, needsToolNote?.let { "$json\n\n$it" } ?: json)
        if (!approved) {
            log.info("freestyle: execution plan rejected by operator for session={}", sessionId.value)
            emitRejected(sessionId, "rejected by operator at plan review", "operator")
            return
        }
        eventStore.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = Clock.System.now(),
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = ExecutionPlanLockedEvent(
                    sessionId = sessionId,
                    planArtifactId = ArtifactId("execution_plan"),
                    workflowId = graph.id,
                    stageIds = graph.stageIds.map { it.value },
                    startStageId = graph.start.value,
                ),
            ),
        )
        runPhase2(sessionId, graph, config)
    }

    /**
     * Recompiles the stored execution plan for resume-after-restart. Requires the
     * orchestrator's artifact cache to be rehydrated first. Null when the plan is
     * absent or no longer compiles — callers fall back to a route-level error.
     */
    fun compiledGraph(sessionId: SessionId): WorkflowGraph? {
        val json = planContent(sessionId) ?: return null
        return runCatching { compiler.compile(json, "freestyle-${sessionId.value}", sessionArtifacts(sessionId)) }
            .onFailure { log.error("freestyle: plan recompile for resume failed: {}", it.message) }
            .getOrNull()
    }

    /** Artifact ids validated earlier in the session (e.g. planning-phase `dod`) for the #264 needs seam. */
    private fun sessionArtifacts(sessionId: SessionId): Set<String> =
        eventStore.read(sessionId)
            .mapNotNull { (it.payload as? ArtifactValidatedEvent)?.artifactId?.value }
            .toSet()

    /**
     * Grounds [graph] against recorded workspace facts and emits [PlanGroundingEvaluatedEvent].
     * Returns null when the plan grounds (PASS), else the findings summary — the caller decides
     * whether to return the plan to architect or reject it. Missing repo map/profile (fresh session,
     * no scan) grounds vacuously; the manifest-produced-by-plan check still catches "build with
     * nothing to build".
     */
    private suspend fun groundPlan(sessionId: SessionId, graph: WorkflowGraph): String? {
        val events = eventStore.read(sessionId)
        val repoMap = events.mapNotNull { it.payload as? RepoMapComputedEvent }.lastOrNull()
        val profile = events.mapNotNull { it.payload as? ProjectProfileBoundEvent }.lastOrNull()
        val paths = repoMap?.entries?.map { it.path }?.toSet().orEmpty()
        // No RepoMapComputedEvent = no scan ran, so `paths` is unknown, not "empty workspace".
        // Tell the grounder not to prove paths absent from a set it never observed (the false
        // "apps/server/** doesn't exist" reject); the build-manifest check still runs.
        val result = PlanGrounder.ground(graph, paths, profile?.commands.orEmpty(), scanned = repoMap != null)
        eventStore.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = Clock.System.now(),
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = PlanGroundingEvaluatedEvent(
                    sessionId = sessionId,
                    planId = graph.id,
                    stateKey = repoMap?.stateKey.orEmpty(),
                    verdict = result.verdict,
                    findings = result.findings,
                ),
            ),
        )
        if (result.verdict == PlanGroundingVerdict.PASS) return null
        return result.findings.joinToString("; ")
    }

    private suspend fun emitPlanLint(sessionId: SessionId, candidateId: String, lint: PlanLintResult) {
        eventStore.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = Clock.System.now(),
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = PlanLintCompletedEvent(
                    sessionId = sessionId,
                    candidateId = candidateId,
                    hardFailures = lint.hardFailures,
                    softFindings = lint.softFindings,
                    score = lint.score,
                ),
            ),
        )
    }

    private suspend fun emitCapabilityGap(
        sessionId: SessionId,
        stageId: String,
        impliedCapability: String,
        evidence: String,
    ) {
        eventStore.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = Clock.System.now(),
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = CapabilityGapDetectedEvent(
                    sessionId = sessionId,
                    stageId = StageId(stageId),
                    impliedCapability = impliedCapability,
                    evidence = evidence,
                    timestampMs = Clock.System.now().toEpochMilliseconds(),
                ),
            ),
        )
    }

    /**
     * Runs the bounded reflection pass (if wired) over [gaps], emits one
     * [CapabilityGapReflectedEvent] per gap, and returns an operator-facing note listing any
     * NEEDS_TOOL verdicts — null when there are no gaps, no reflector, or nothing to escalate.
     * A single call for the whole batch: no per-gap loop, no retry (bounded per invariant #9).
     */
    private suspend fun reflectOnGaps(sessionId: SessionId, gaps: List<CapabilityGap>): String? {
        val judge = reflector
        if (gaps.isEmpty() || judge == null) return null
        val inputs = gaps.map {
            CapabilityGapInput(it.stageId, it.impliedCapability, it.phrase, it.grantedCapabilities)
        }
        val reflections = runCatching { judge.reflect(sessionId, inputs) }.getOrElse {
            log.warn("[FreestyleDriver] capability gap reflection failed session={}: {}", sessionId.value, it.message)
            emptyList()
        }
        reflections.forEach { r ->
            log.info(
                "[FreestyleDriver] capability gap reflected session={} stage={} implied={} verdict={}",
                sessionId.value, r.stageId, r.impliedCapability, r.verdict,
            )
            emitCapabilityGapReflected(sessionId, r.stageId, r.impliedCapability.name, r.verdict, r.detail)
        }
        val escalations = reflections.filter { it.verdict == CapabilityGapVerdict.NEEDS_TOOL }
        return escalations.takeIf { it.isNotEmpty() }?.let {
            "CAPABILITY GAPS REQUESTED (reflection confirmed a genuine need — approve or reject the plan " +
                "with this in mind; no tool is auto-granted):\n" +
                it.joinToString("\n") { e -> "- stage ${e.stageId} needs ${e.impliedCapability}: ${e.detail}" }
        }
    }

    private suspend fun emitCapabilityGapReflected(
        sessionId: SessionId,
        stageId: String,
        impliedCapability: String,
        verdict: CapabilityGapVerdict,
        detail: String,
    ) {
        eventStore.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = Clock.System.now(),
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = CapabilityGapReflectedEvent(
                    sessionId = sessionId,
                    stageId = StageId(stageId),
                    impliedCapability = impliedCapability,
                    verdict = verdict,
                    detail = detail,
                    timestampMs = Clock.System.now().toEpochMilliseconds(),
                ),
            ),
        )
    }

    private suspend fun emitRejected(sessionId: SessionId, reason: String, source: String) {
        eventStore.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = Clock.System.now(),
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = ExecutionPlanRejectedEvent(
                    sessionId = sessionId,
                    reason = reason,
                    source = source,
                ),
            ),
        )
        // A rejected plan ends the session — but the last workflow verdict on record was the
        // planning phase's WorkflowCompleted, so the session read as SUCCESS (the "COMPLETED-lie").
        // Emit a session-terminal WorkflowFailed so the run's true outcome is on the log. stageId is
        // architect: the plan's producer and where a fix (or future return-to-architect loop) lands.
        eventStore.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = Clock.System.now(),
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = WorkflowFailedEvent(
                    sessionId = sessionId,
                    stageId = StageId("architect"),
                    reason = "execution plan rejected ($source): $reason",
                    retryExhausted = false,
                ),
            ),
        )
    }

    companion object {
        private val log = LoggerFactory.getLogger(FreestyleDriver::class.java)
    }
}
