package com.correx.core.kernel.orchestration

import com.correx.core.approvals.ApprovalOutcome
import com.correx.core.artifacts.kind.JsonSchema
import com.correx.core.context.model.ContextEntry
import com.correx.core.context.model.ContextLayer
import com.correx.core.context.model.EntryRole
import com.correx.core.events.events.ApprovalDecisionResolvedEvent
import com.correx.core.events.events.ApprovalRequestedEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.events.ClarificationAnswer
import com.correx.core.events.events.ClarificationAnsweredEvent
import com.correx.core.events.events.ClarificationQuestions
import com.correx.core.events.events.ClarificationRequestedEvent
import com.correx.core.events.events.RepoKnowledgeHit
import com.correx.core.events.events.RepoKnowledgeRetrievedEvent
import com.correx.core.events.events.RepoMapComputedEvent
import com.correx.core.events.events.OrchestrationPausedEvent
import com.correx.core.events.events.OrchestrationResumedEvent
import com.correx.core.events.events.ExecutionPlanLockedEvent
import com.correx.core.events.events.InitialIntentEvent
import com.correx.core.events.events.SteeringNoteAddedEvent
import com.correx.core.events.types.ArtifactId
import com.correx.core.events.types.ClarificationRequestId
import com.correx.core.events.types.ContextEntryId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.inference.ResponseFormat
import com.correx.core.transitions.evaluation.EvaluationContext
import com.correx.core.transitions.graph.StageConfig
import com.correx.core.transitions.graph.WorkflowGraph
import kotlinx.coroutines.CompletableDeferred
import kotlinx.serialization.encodeToString
import kotlinx.serialization.json.Json
import java.util.*
import java.util.concurrent.*
import java.util.concurrent.atomic.*
import kotlin.coroutines.cancellation.CancellationException

/** Cache mapping "sessionId:artifactId" to JSON-serialized artifact payload content.
 *  Populated when ProcessResult artifacts are stored on tool failure.
 *  Used by DefaultSessionOrchestrator.step() to populate EvaluationContext.artifactContent. */

internal fun SessionOrchestrator.evictArtifactContentCache(sessionId: SessionId) {
    val prefix = "${sessionId.value}:"
    artifactContentCache.keys.removeAll { it.startsWith(prefix) }
}

/** Deterministic extraction/repair ladder for near-miss LLM artifact text (prose-wrapped JSON,
 *  code fences, trailing commas). Pure — recomputes on replay, records no events. */

internal suspend fun SessionOrchestrator.buildSchemaEntries(
    responseFormat: ResponseFormat,
    stageId: StageId,
): List<ContextEntry> {
    if (responseFormat !is ResponseFormat.Json) return emptyList()
    val compactSchema = Json.encodeToString(JsonSchema.serializer(), responseFormat.schema)
    val instruction = "Respond with a single JSON object matching this schema. " +
            "Do not include markdown, code fences, or commentary outside the JSON. " +
            "Schema: $compactSchema"
    return listOf(
        ContextEntry(
            id = ContextEntryId(UUID.randomUUID().toString()),
            layer = ContextLayer.L0,
            content = instruction,
            sourceType = "schemaInstruction",
            sourceId = stageId.value,
            tokenEstimate = estimateTokens(instruction),
            role = EntryRole.SYSTEM,
        ),
    )
}

internal suspend fun SessionOrchestrator.buildSteeringNoteEntries(sessionId: SessionId): List<ContextEntry> {
    val events = eventStore.read(sessionId)
    var locked = false
    return events.mapNotNull { event ->
        when (val p = event.payload) {
            is ExecutionPlanLockedEvent -> { locked = true; null }
            is SteeringNoteAddedEvent -> ContextEntry(
                id = ContextEntryId(UUID.randomUUID().toString()),
                layer = if (locked) ContextLayer.L0 else ContextLayer.L2,
                content = if (locked) "PRIORITY DIRECTIVE — overrides the locked plan: ${p.content}"
                          else p.content,
                sourceType = "steeringNote",
                sourceId = p.stageId?.value ?: sessionId.value,
                tokenEstimate = estimateTokens(p.content),
                // #312: an operator steering note arrives mid-run by definition — mutable ⇒ USER,
                // same as the unlocked path below. Locked notes keep L0 so they stay pinned against
                // the budget; PromptRenderer still restates every note as a trailing anchor.
                role = EntryRole.USER,
            )
            // An operator steering note attached to a decision is a real instruction; keep it.
            // Bare rejections are consolidated separately (buildRejectionFeedbackEntry) so they
            // carry which call/why instead of a repeated context-free warning.
            is ApprovalDecisionResolvedEvent -> p.userSteering?.let { steering ->
                ContextEntry(
                    id = ContextEntryId(UUID.randomUUID().toString()),
                    layer = ContextLayer.L2,
                    content = steering.text,
                    sourceType = "steeringNote",
                    sourceId = steering.sessionId.value,
                    tokenEstimate = estimateTokens(steering.text),
                    role = EntryRole.USER,
                )
            }
            else -> null
        }
    }
}

/**
 * One consolidated, stage-scoped entry naming the tool calls the operator declined in THIS stage —
 * tool, args preview, tier, and the rejection reason — joined from each rejected
 * [ApprovalDecisionResolvedEvent] back to its [ApprovalRequestedEvent] by requestId. Replaces the
 * old per-rejection generic warnings ("a previous tool call was rejected"), which repeated without
 * saying which call or why and read to the model as noise rather than a correction signal. Empty
 * when the stage has no rejected calls. Replay-safe (recorded events only).
 */
fun buildRejectionFeedbackEntry(events: List<StoredEvent>, stageId: StageId): ContextEntry? {
    val requests = events.mapNotNull { it.payload as? ApprovalRequestedEvent }.associateBy { it.requestId }
    val rejections = events.mapNotNull { it.payload as? ApprovalDecisionResolvedEvent }
        .filter { it.outcome == ApprovalOutcome.REJECTED && requests[it.requestId]?.stageId == stageId }
    if (rejections.isEmpty()) return null
    val content = buildString {
        appendLine("## Rejected tool calls this stage")
        appendLine(
            "The operator declined the calls below. Do not re-propose them as-is — address the " +
                "stated reason, or take a different approach:",
        )
        rejections.forEach { d ->
            val req = requests[d.requestId]
            val preview = req?.preview?.takeIf { it.isNotBlank() }?.let { " $it" }.orEmpty()
            val reason = d.reason?.takeIf { it.isNotBlank() }
                ?.let { " — reason: $it" } ?: " — no reason given"
            appendLine("- ${req?.toolName ?: "tool call"} (tier ${d.tier})$preview$reason")
        }
    }.trimEnd()
    return ContextEntry(
        id = ContextEntryId(UUID.randomUUID().toString()),
        layer = ContextLayer.L2,
        content = content,
        sourceType = "rejectionFeedback",
        sourceId = stageId.value,
        tokenEstimate = content.length / 4,
        // #312: literally operator voice ("the operator declined"), and it grows mid-stage as more
        // calls are rejected — USER, trailing slot.
        role = EntryRole.USER,
    )
}

/**
 * Injects the initial user intent (the freeform request that started the run) as an L1 USER entry
 * present in EVERY stage's context (architecture-conformance, 2026-07-14). The intent
 * is the single most load-bearing constraint of a run, yet it previously reached normal stages only
 * as a repo-map retrieval seed (repoKnowledgeQuery) or, on the rare Tier-2 recovery path, the
 * arbiter ticket — so an implementer could drift from the goal with the goal itself absent from its
 * authoritative context. It is REQUIRED-bucket, so it is never dropped under budget. Absent for
 * fixed-task workflows (no InitialIntentEvent) → empty.
 */

internal suspend fun SessionOrchestrator.buildIntentEntry(sessionId: SessionId): List<ContextEntry> {
    val intent = initialIntent(sessionId)?.takeIf { it.isNotBlank() } ?: return emptyList()
    val content = "## Authoritative intent\nThe operator's request that started this run. Everything " +
        "you produce must serve it; do not drift from it or expand its scope:\n$intent"
    return listOf(
        ContextEntry(
            id = ContextEntryId(UUID.randomUUID().toString()),
            // #290: the original intent is the operator's request, not Correx policy — render it as
            // L1/USER (not folded into leading system). It stays unprunable via the REQUIRED bucket
            // (initialIntent ∈ REQUIRED_SOURCE_TYPES), so retention no longer rides on L0/SYSTEM.
            layer = ContextLayer.L1,
            content = content,
            sourceType = "initialIntent",
            sourceId = sessionId.value,
            tokenEstimate = estimateTokens(content),
            role = EntryRole.USER,
        ),
    )
}

/** The freeform request that started the run, or null for a fixed-task workflow. */
internal fun SessionOrchestrator.initialIntent(sessionId: SessionId): String? =
    eventStore.read(sessionId)
        .mapNotNull { it.payload as? InitialIntentEvent }
        .firstOrNull()?.intent?.trim()

/**
 * Injects the operator's answers to a stage's open questions as a pinned L0 SYSTEM entry, so the
 * stage sees its own questions resolved on the clarification re-run. The prompt calls these answers
 * "authoritative", so they are pinned as standing context (L0) rather than a droppable L2 turn —
 * matching the mechanics to the stated authority. They render as USER (#312): these are the
 * operator's own words, and the set grows with each clarification round. Correlates each answer's
 * questionId back to the prompt recorded on the [ClarificationRequestedEvent].
 */

internal suspend fun SessionOrchestrator.buildClarificationAnswerEntries(sessionId: SessionId): List<ContextEntry> {
    val events = eventStore.read(sessionId)
    val answered = events.mapNotNull { it.payload as? ClarificationAnsweredEvent }
    if (answered.isEmpty()) return emptyList()
    val prompts = events
        .mapNotNull { it.payload as? ClarificationRequestedEvent }
        .flatMap { it.questions }
        .associate { it.id to it.prompt }
    val rendered = answered.flatMap { it.answers }.joinToString("\n") { a ->
        "Q: ${prompts[a.questionId] ?: a.questionId}\nA: ${a.value}"
    }
    val content = "The operator answered your open questions. Treat these answers as authoritative " +
        "and do not ask them again:\n$rendered"
    return listOf(
        ContextEntry(
            id = ContextEntryId(UUID.randomUUID().toString()),
            layer = ContextLayer.L0,
            content = content,
            sourceType = "clarificationAnswer",
            sourceId = sessionId.value,
            tokenEstimate = estimateTokens(content),
            // #312: operator's own words, and the set grows per clarification round — USER.
            role = EntryRole.USER,
        ),
    )
}

/**
 * If [stageId] just produced an LLM artifact carrying a non-empty `questions` array, park the
 * run: record the questions (invariant #9 — observed at parse time), await the operator's
 * answers, and record them. The caller then ADVANCES to the next stage — the stage is not
 * re-run. The answers are session-wide events, injected into every later stage's L0 context
 * (see [buildClarificationAnswerEntries]), so the next stage acts on them directly; re-running
 * the questioning stage only restarts inference with no prior reasoning and makes it re-explore.
 * Bounded by [OrchestrationTuning.maxClarificationRounds] so retries can't re-park unboundedly.
 * Returns true when it parked and collected answers, false when there are no questions or the
 * round budget is spent.
 */

internal suspend fun SessionOrchestrator.requestClarificationIfNeeded(
    sessionId: SessionId,
    stageId: StageId,
    graph: WorkflowGraph,
): Boolean {
    val slot = graph.stages[stageId]?.produces?.firstOrNull { it.kind.llmEmitted } ?: return false
    val content = artifactContentCache["${sessionId.value}:${slot.name.value}"] ?: return false
    val questions = ClarificationQuestions.parse(content)
    if (questions.isEmpty()) return false

    val priorRounds = eventStore.read(sessionId)
        .count { (it.payload as? ClarificationRequestedEvent)?.stageId == stageId }
    if (priorRounds >= tuning.maxClarificationRounds) return false

    val requestId = ClarificationRequestId(UUID.randomUUID().toString())
    val deferred = CompletableDeferred<List<ClarificationAnswer>>()
    pendingClarifications[requestId] = deferred
    emit(sessionId, OrchestrationPausedEvent(sessionId, stageId, "CLARIFICATION_PENDING"))
    emit(sessionId, ClarificationRequestedEvent(requestId, sessionId, stageId, questions))
    return try {
        val answers = deferred.await()
        emit(sessionId, ClarificationAnsweredEvent(requestId, sessionId, stageId, answers))
        emit(sessionId, OrchestrationResumedEvent(sessionId, stageId))
        true
    } finally {
        pendingClarifications.remove(requestId)
    }
}

/**
 * Renders the recorded repo map as a structural "what exists" listing — directories with
 * their files, alphabetical, no recency ranking and no symbol excerpts. Relevance is the
 * retrieval layer's job (buildRelevantFilesEntry); the map only tells the model what it can
 * `file_read`. Markdown/docs paths are gated: they only appear when the stage prompt
 * explicitly asks about documentation, so churny kernel docs can't bury real source
 * (2026-07-05 context-poisoning root cause). Reads the latest [RepoMapComputedEvent] from
 * the log (replay-safe — never re-scans the FS).
 */

internal suspend fun SessionOrchestrator.buildRepoMapEntries(sessionId: SessionId, stagePrompt: String = ""): List<ContextEntry> {
    val map = eventStore.read(sessionId)
        .mapNotNull { it.payload as? RepoMapComputedEvent }
        .lastOrNull() ?: return emptyList()
    val docsWanted = DOCS_REQUEST_REGEX.containsMatchIn(stagePrompt)
    val paths = map.entries.map { it.path }.distinct()
        .filter { docsWanted || !isDocPath(it) }
        .sorted()
    if (paths.isEmpty()) return emptyList()
    val byDir = paths.groupBy { it.substringBeforeLast('/', missingDelimiterValue = ".") }
    val content = buildString {
        appendLine("## Repo layout (what exists — use file_read for content)")
        byDir.entries.sortedBy { it.key }.forEach { (dir, files) ->
            val names = files.map { it.substringAfterLast('/') }
            val shown = names.take(tuning.repoMapFilesPerDir)
            val more = names.size - shown.size
            val suffix = if (more > 0) ", …(+$more more)" else ""
            appendLine("- $dir/ (${names.size}): ${shown.joinToString(", ")}$suffix")
        }
    }.trimEnd()
    return listOf(
        ContextEntry(
            id = ContextEntryId(UUID.randomUUID().toString()),
            layer = ContextLayer.L3,
            content = content,
            sourceType = "repoMap",
            sourceId = "repo-map",
            tokenEstimate = estimateTokens(content),
            // #290: L3 retrieval is reference, not policy — USER role so it stays out of leading system.
            role = EntryRole.USER,
        ),
    )
}

/**
 * True when [stageId]'s own prompt must be suppressed because the stage was entered to repair a
 * gate failure. That happens when the most recent transition INTO this stage was a ticket route
 * ([TICKET_ROUTE], same entry check as DefaultSessionOrchestrator.ticketReturnMove) — the failure
 * ticket routes an owner stage back to fix a file it authored, and its original generative mandate
 * ("scaffold the project") is exactly what drives it to overwrite that file with a stub. The
 * recovery-ticket entry ([buildRecoveryTicketEntry]) becomes the sole mandate instead.
 *
 * The purpose-built recovery arbiter (role=recovery) is excluded: it is also ticket-entered but
 * its own prompt IS a tuned repair prompt, so it keeps it. Replay-safe (recorded events only).
 */

internal suspend fun SessionOrchestrator.buildContextualRepoEntries(
    sessionId: SessionId,
    stageId: StageId,
    stageConfig: StageConfig,
): List<ContextEntry> {
    val stagePrompt = repoKnowledgeQuery(sessionId, stageId, stageConfig)
    val recorded = eventStore.read(sessionId)
        .mapNotNull { it.payload as? RepoKnowledgeRetrievedEvent }
        .lastOrNull { it.stageId == stageId }
    if (recorded != null) {
        return repoEntriesOrMapFloor(sessionId, recorded.hits, stagePrompt) + buildDocsCatalogEntry(sessionId)
    }

    // Sibling-stage output written earlier this session isn't in the session-start repo map and was
    // never embedded, so semantic retrieval is structurally blind to it. Overlay it as deterministic
    // hits (score 1.0) that lead the semantic hits, deduped by path — no reindex, no embed.
    val writtenHits = sessionWrittenHits(sessionId, stageId)
    val retriever = repoKnowledgeRetriever
    val semanticHits = if (retriever == null) {
        emptyList()
    } else {
        runCatching { retriever.retrieve(sessionId, stagePrompt, tuning.repoMapInjectTopK) }
            .getOrElse { e ->
                if (e is CancellationException) throw e
                log.warn("repo-knowledge retrieval failed for stage {}: {}", stageId.value, e.message)
                emptyList()
            }
    }
    val writtenPaths = writtenHits.mapTo(mutableSetOf()) { it.path }
    val hits = writtenHits + semanticHits.filterNot { it.path in writtenPaths }
    // Record the raw retrieval for audit (a degenerate all-zero result is itself the signal that
    // the embedder is noop / the index is empty), but only useful hits ever reach the agent.
    if (hits.isNotEmpty()) emit(sessionId, RepoKnowledgeRetrievedEvent(sessionId, stageId, stagePrompt, hits))
    return repoEntriesOrMapFloor(sessionId, hits, stagePrompt) + buildDocsCatalogEntry(sessionId)
}

/**
 * Always-on, capped catalog of the docs that exist, as `path — descriptor` lines (the descriptor
 * is the doc's frontmatter/H1/first-line summary recorded on [RepoMapEntry.symbols]). This is the
 * read-on-demand doc index: one line per doc lets a stage decide whether to file_read it, so docs
 * stay discoverable without the heading-flood the docs gate was added to suppress. Distinct from
 * [buildRepoMapEntries]/[repoEntriesOrMapFloor], which still exclude doc *content* from the source
 * layout for non-doc prompts. Reads the recorded map (replay-safe — never re-scans the FS).
 */

internal suspend fun SessionOrchestrator.buildDocsCatalogEntry(sessionId: SessionId): List<ContextEntry> {
    val map = eventStore.read(sessionId)
        .mapNotNull { it.payload as? RepoMapComputedEvent }
        .lastOrNull() ?: return emptyList()
    val docs = map.entries
        .filter { isDocPath(it.path) }
        .sortedByDescending { it.score }
        .take(tuning.docsCatalogMax)
    if (docs.isEmpty()) return emptyList()
    val content = buildString {
        appendLine("## Docs available (file_read to open — do not assume contents)")
        docs.forEach { entry ->
            val descriptor = entry.symbols.firstOrNull()?.takeIf { it.isNotBlank() }
            appendLine(if (descriptor != null) "- ${entry.path} — $descriptor" else "- ${entry.path}")
        }
    }.trimEnd()
    return listOf(
        ContextEntry(
            id = ContextEntryId(UUID.randomUUID().toString()),
            layer = ContextLayer.L3,
            content = content,
            sourceType = "docsCatalog",
            sourceId = "docs-catalog",
            tokenEstimate = estimateTokens(content),
            // #290: docs catalog is L3 reference — USER role so it stays out of leading system.
            role = EntryRole.USER,
        ),
    )
}

/**
 * Zero-score hits carry no ranking signal (a noop embedder scores everything 0.0), so injecting
 * them just feeds the agent arbitrary files and invites hallucinated paths. Drop them; when nothing
 * useful survives, fall back to the deterministic repo map so the stage still has real grounding.
 *
 * Markdown/docs hits are gated exactly like the repo-map floor (buildRepoMapEntries): with the
 * embedder live, semantic similarity happily ranks churny kernel docs (epics, ADRs, QA plans)
 * above real source for a code task, re-poisoning the context the floor's gate was added to fix
 * (2026-07-05 root cause). They only survive when the stage prompt explicitly asks about docs.
 */

internal suspend fun SessionOrchestrator.repoEntriesOrMapFloor(
    sessionId: SessionId,
    hits: List<RepoKnowledgeHit>,
    stagePrompt: String,
): List<ContextEntry> {
    val docsWanted = DOCS_REQUEST_REGEX.containsMatchIn(stagePrompt)
    val useful = hits.filter { it.score > 0f && (docsWanted || !isDocPath(it.path)) }
    return if (useful.isEmpty()) {
        buildRepoMapEntries(sessionId, stagePrompt)
    } else {
        listOf(buildRelevantFilesEntry(useful))
    }
}

/**
 * Stamps the required/optional budget bucket by sourceType. REQUIRED = the stage cannot run
 * correctly without it (task, constraints, loaded inputs, failure feedback) — never trimmed,
 * build fails fast if it alone exceeds the budget. Everything else (repo map, retrieval,
 * journal, profiles, vocabulary, transcript) is OPTIONAL and shares the remainder under
 * per-type ceilings.
 */

internal suspend fun SessionOrchestrator.buildNeedsArtifactEntries(
    sessionId: SessionId,
    stageConfig: StageConfig,
    criticIds: Set<ArtifactId> = emptySet(),
    criticFromStage: String? = null,
): List<ContextEntry> {
    return stageConfig.needs.mapNotNull { needed ->
        val cached = artifactContentCache["${sessionId.value}:${needed.value}"]
            ?.takeIf { it.isNotBlank() } ?: return@mapNotNull null
        val content = if (needed in criticIds) {
            "## Critic feedback from stage '${criticFromStage ?: "reviewer"}': ${needed.value}\n" +
                "This is reviewer feedback on your previous attempt. " +
                "Address every finding before completing the stage.\n$cached"
        } else {
            // file_written handoff: hand the next stage the whole producing stage's change set
            // as a manifest (known files), not one file's content. The cache value only holds
            // the LAST write of the stage, so rendering it inline both misleads (5 of 6 files
            // invisible) and wastes budget — content stays lazy behind file_read.
            val manifest = parseFileWrittenArtifact(cached)
                ?.let { fileWrittenManifest(sessionId, needed) }
            "## Input artifact: ${needed.value}\n${manifest ?: cached}"
        }
        ContextEntry(
            id = ContextEntryId(UUID.randomUUID().toString()),
            layer = ContextLayer.L1,
            content = content,
            sourceType = if (needed in criticIds) "criticFeedback" else "neededArtifact",
            sourceId = needed.value,
            tokenEstimate = estimateTokens(content),
            role = EntryRole.USER,
        )
    }
}
