package com.correx.core.context.builder

import com.correx.core.context.compression.CompressionPolicy
import com.correx.core.context.compression.CompressionStage
import com.correx.core.context.compression.CompressionStrategy
import com.correx.core.context.compression.ContextClass
import com.correx.core.context.compression.ContextClassifier
import com.correx.core.context.compression.ContextCompressor
import com.correx.core.context.compression.FactSheet
import com.correx.core.context.compression.FormatCompressor
import com.correx.core.context.compression.NoOpTokenPruner
import com.correx.core.context.compression.ProtectedSpanTagger
import com.correx.core.context.compression.RelevanceScorer
import com.correx.core.context.compression.ToMeMerger
import com.correx.core.context.compression.TokenPruner
import com.correx.core.context.compression.UniformRelevanceScorer
import com.correx.core.context.model.CompressionMetadata
import com.correx.core.context.model.ContextBucket
import com.correx.core.context.model.ContextEntry
import com.correx.core.context.model.ContextLayer
import com.correx.core.context.model.ContextPack
import com.correx.core.context.model.EntryRole
import com.correx.core.context.model.OptionalCeilings
import com.correx.core.context.model.TokenBudget
import com.correx.core.events.types.ContextEntryId
import com.correx.core.events.types.ContextPackId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.jsonPrimitive

/**
 * Runs the compression pipeline (docs/plans/correx-compression-pipeline.md) in fixed pipeline
 * order; each stage is skipped unless its [CompressionStage] is enabled by the [policy] level.
 * Lossy stages (token pruning, relevance reordering) are suspend because they may hit an
 * out-of-process sidecar; the deterministic stages run in-line. Recursive summarization
 * (level 5) is handled by the orchestrator's TierContextSummarizer, not here, because it emits
 * events — the builder stays free of side effects.
 */
class DefaultContextPackBuilder(
    private val compressor: ContextCompressor,
    private val policy: CompressionPolicy = CompressionPolicy(),
    private val classifier: ContextClassifier = ContextClassifier(),
    private val formatCompressor: FormatCompressor = FormatCompressor(),
    private val tagger: ProtectedSpanTagger = ProtectedSpanTagger(),
    private val factSheet: FactSheet = FactSheet(ProtectedSpanTagger()),
    private val tokenPruner: TokenPruner = NoOpTokenPruner,
    private val toMeMerger: ToMeMerger = ToMeMerger(),
    private val relevanceScorer: RelevanceScorer = UniformRelevanceScorer,
    private val ceilings: OptionalCeilings = OptionalCeilings(),
) : ContextPackBuilder {

    // decisionJournal is deliberately NOT here anymore: the journal is OPTIONAL-bucket, capped
    // by OptionalCeilings. Pinning is otherwise bucket-driven (ContextBucket.REQUIRED); this set
    // only covers entries built outside the orchestrator that never stamp buckets.
    // "remainingDelta" is the shrinking stage-contract checklist (stage-termination design
    // 2026-07-11): it must survive every budget/dedup pass, since its entire purpose is to give
    // the model a progress signal that outlives the truncation which otherwise wipes its memory.
    // "retryFeedback" carries WHY the last attempt failed + the files already written this stage; if
    // truncation evicts it the model cold-starts on the same wrong idea every turn (the write-loop rot).
    private val neverDropSourceTypes = setOf("steeringNote", "eventHistory", "factSheet", "remainingDelta", "retryFeedback")

    private companion object {
        const val CHARS_PER_TOKEN = 4
        const val FACT_SHEET_ORDINAL = -1
        const val DEFAULT_PRUNE_RATIO = 0.45
        const val DOC_PRUNE_RATIO = 0.6
        const val DOC_MIN_CHARS = 1200
        const val TIER0_TURNS = 3
        // Instruction-doc LLMLingua pruning is retired: it fused load-bearing procedural text
        // (the numbered stage_complete / emit_artifact steps in AGENTS.md, the curated
        // project.toml profile) into unparseable soup, and the agent could no longer tell how to
        // finish a stage. The static block is instead kept small by DROPPING CLAUDE.md at the
        // loader — not by garbling what remains. Left empty (not deleted) so the branch is a
        // no-op guard rather than dead-code churn; freeform conversation pruning is unaffected.
        val DOC_SOURCE_TYPES = emptySet<String>()
    }

    override suspend fun build(
        id: ContextPackId,
        sessionId: SessionId,
        stageId: StageId,
        entries: List<ContextEntry>,
        budget: TokenBudget
    ): ContextPack {
        // Stamp chronological order from the caller's input sequence so grouping/compression can
        // reorder freely and PromptRenderer can restore true turn order from the ordinal.
        val stamped = entries.mapIndexed { index, entry -> entry.copy(ordinal = index) }

        // A stalled stage re-issuing the exact same call otherwise burns budget on N identical
        // (tool, arguments) pairs that all say the same thing — collapse those to the most recent
        // occurrence. Fingerprint is (name, arguments) verbatim, so a call whose arguments genuinely
        // differ (different start_line/end_line, different path) is a distinct read and untouched.
        // Must run before FORMAT_COMPRESS below: that stage reencodes assistantToolCall content from
        // raw ToolCallRequest JSON into a dotted-line format the fingerprint parser can't read.
        val deduped = dedupeRepeatedToolCalls(stamped)

        // #289: a successful file_write/file_edit echoes the whole file body (+ diff) back in its
        // tool result — up to ~30k chars, which pushed the traced Gemma4 request past its context
        // window. The write already happened and its full output is durable in the event log/CAS;
        // the model only needs a receipt that it landed. Compact those results to a one-line receipt
        // (reads and gate output stay verbatim). Derived-only — authoritative events are untouched.
        val compacted = compactWriteReceipts(deduped)

        // Stage 1 (FORMAT_COMPRESS): lossless json→dotted-line compaction of structured entries.
        val formatted = if (policy.enabled(CompressionStage.FORMAT_COMPRESS)) {
            compacted.map { entry ->
                if (classifier.classify(entry) == ContextClass.STRUCTURED) {
                    reencode(entry, formatCompressor.compress(entry.content))
                } else {
                    entry
                }
            }
        } else compacted

        // Stage 3 (TOKEN_PRUNE): prune freeform prose, preserving protected spans. When TIER_SPLIT
        // is on, the newest TIER0_TURNS freeform turns are left full-fidelity (tier 0).
        val pruned = if (policy.enabled(CompressionStage.TOKEN_PRUNE)) prune(formatted) else formatted

        // Stage 8 (TOME_MERGE): collapse near-duplicate freeform turns into the newest survivor.
        val merged = if (policy.enabled(CompressionStage.TOME_MERGE)) applyToMe(pruned) else pruned

        // Stage 7 (FACT_SHEET): pin load-bearing spans as an L0 SYSTEM entry that never drops.
        val withFactSheet = if (policy.enabled(CompressionStage.FACT_SHEET)) {
            val facts = factSheet.extract(merged)
            if (facts.isEmpty()) merged else merged + factSheetEntry(factSheet.render(facts))
        } else merged

        val (pinned, compressible) = withFactSheet.partition {
            it.bucket == ContextBucket.REQUIRED || it.sourceType in neverDropSourceTypes
        }
        // Required entries are never trimmed: if they alone blow the budget, fail fast rather
        // than ship a silently truncated instruction set.
        val required = withFactSheet.filter { it.bucket == ContextBucket.REQUIRED }
        val requiredTokens = required.sumOf { it.tokenEstimate }
        if (requiredTokens > budget.limit) {
            throw RequiredContextOverflowException(
                requiredTokens = requiredTokens,
                budgetLimit = budget.limit,
                breakdown = required.groupBy { it.sourceType }
                    .mapValues { (_, group) -> group.sumOf { it.tokenEstimate } },
            )
        }
        val pinnedTokens = pinned.sumOf { it.tokenEstimate }
        var remainingTokens = (budget.limit - pinnedTokens).coerceAtLeast(0)

        // Optional per-type ceilings are enforced by truncation here, not left to the group
        // compressor: several strategies (EventHistory) pass entries through untouched, and these
        // feeds are single large entries — dropping whole entries would lose them entirely.
        val ceilinged = compressible.map(::applyCeiling)

        // Query-conditioned selection: when a real relevance scorer is wired, reorder each freeform
        // group so the least-relevant turns sit at the front (dropped first) and the most-relevant
        // land at the tail the Conversation strategy keeps. Final layering re-sorts by ordinal, so
        // chronological render order is unaffected — only *what survives* the budget changes.
        val relevanceRanked = rankByRelevance(ceilinged, currentQuery(withFactSheet))

        // Stage 9 (BUDGET_ALLOCATOR_STRICT): compress structured groups before freeform so a long
        // conversation can't starve tool logs / artifacts out of the window.
        val strict = policy.enabled(CompressionStage.BUDGET_ALLOCATOR_STRICT)
        val strategiesUsed = linkedSetOf<String>()
        val grouped = relevanceRanked.groupBy { it.sourceType }.entries
        val orderedGroups = if (strict) {
            grouped.sortedBy { (_, group) -> if (classifier.classify(group.first()) == ContextClass.FREEFORM) 1 else 0 }
        } else grouped.toList()
        val compressed = orderedGroups.flatMap { (sourceType, group) ->
            val strategy = strategyFor(sourceType)
            strategiesUsed += strategy::class.simpleName ?: "Unknown"
            // Optional groups run under their per-type ceiling so no single feed (repo map,
            // journal, profile) can crowd out the rest of the optional block.
            val groupLimit = ceilings.capFor(sourceType)
                ?.coerceAtMost(remainingTokens.coerceAtLeast(0))
                ?: remainingTokens.coerceAtLeast(0)
            val result = compressor.compress(group, TokenBudget(limit = groupLimit), strategy)
            remainingTokens -= result.sumOf { it.tokenEstimate }
            result
        }

        val retained = reconcileToolPairs((pinned + compressed).sortedBy { it.ordinal })
        val layers = retained.groupBy { it.layer }
        val droppedCount = entries.size - retained.count { it.sourceType != "factSheet" }
        val truncatedLayers = if (droppedCount > 0) {
            entries.map { it.layer }.distinct().filter { layer ->
                entries.count { it.layer == layer } > retained.count { it.layer == layer }
            }
        } else emptyList()

        return ContextPack(
            id = id,
            sessionId = sessionId,
            stageId = stageId,
            layers = layers,
            budgetUsed = retained.sumOf { it.tokenEstimate },
            budgetLimit = budget.limit,
            compressionMetadata = CompressionMetadata(
                appliedStrategies = strategiesUsed.toList().ifEmpty { listOf("Conversation") },
                truncatedLayers = truncatedLayers,
                entriesDropped = droppedCount
            )
        )
    }

    private suspend fun prune(entries: List<ContextEntry>): List<ContextEntry> {
        val tierSplit = policy.enabled(CompressionStage.TIER_SPLIT)
        val freeformOrdinals = entries.filter { classifier.classify(it) == ContextClass.FREEFORM }
            .map { it.ordinal }.sortedDescending()
        val tier0 = if (tierSplit) freeformOrdinals.take(TIER0_TURNS).toSet() else emptySet()
        return entries.map { entry ->
            // Required entries are exempt from lossy pruning by definition of the bucket.
            if (entry.bucket == ContextBucket.REQUIRED) return@map entry
            val isFreeform = classifier.classify(entry) == ContextClass.FREEFORM && entry.ordinal !in tier0
            val isLargeDoc = entry.sourceType in DOC_SOURCE_TYPES && entry.content.length >= DOC_MIN_CHARS
            when {
                isFreeform -> reencode(entry, tokenPruner.prune(entry.content, spansOf(entry), DEFAULT_PRUNE_RATIO))
                // Static instruction docs: gentler ratio so directives survive the classifier.
                isLargeDoc -> reencode(entry, tokenPruner.prune(entry.content, spansOf(entry), DOC_PRUNE_RATIO))
                else -> entry
            }
        }
    }

    /**
     * Truncates an OPTIONAL entry to its per-sourceType token ceiling. The decision journal is
     * chronological (newest decisions at the tail), so it keeps its tail; every other feed keeps
     * its head (headers/most-structural lines first).
     */
    private fun applyCeiling(entry: ContextEntry): ContextEntry {
        val cap = ceilings.capFor(entry.sourceType) ?: return entry
        if (entry.tokenEstimate <= cap) return entry
        val trimmedTailMarker = "…(older entries trimmed)\n"
        val trimmedHeadMarker = "\n…(trimmed to budget)"
        val content = if (entry.sourceType == "decisionJournal") {
            trimmedTailMarker + entry.content.takeLast(cap * CHARS_PER_TOKEN - trimmedTailMarker.length)
        } else {
            entry.content.take(cap * CHARS_PER_TOKEN - trimmedHeadMarker.length) + trimmedHeadMarker
        }
        return entry.copy(content = content, tokenEstimate = estimateTokens(content))
    }

    private fun spansOf(entry: ContextEntry) = tagger.protectedSpans(entry.content)

    private fun applyToMe(entries: List<ContextEntry>): List<ContextEntry> {
        val (freeform, rest) = entries.partition {
            it.bucket != ContextBucket.REQUIRED && classifier.classify(it) == ContextClass.FREEFORM
        }
        return (rest + toMeMerger.merge(freeform)).sortedBy { it.ordinal }
    }

    // The current query is the most-recent user turn — the anchor relevance is scored against.
    private fun currentQuery(entries: List<ContextEntry>): String? =
        entries.filter { it.role == EntryRole.USER }.maxByOrNull { it.ordinal }?.content

    private suspend fun rankByRelevance(entries: List<ContextEntry>, query: String?): List<ContextEntry> {
        if (relevanceScorer is UniformRelevanceScorer || query == null) return entries
        val (freeform, rest) = entries.partition { classifier.classify(it) == ContextClass.FREEFORM }
        if (freeform.isEmpty()) return entries
        val scores = relevanceScorer.score(query, freeform.map { it.content })
        // Least-relevant first so the Conversation strategy's takeLast keeps the most-relevant.
        val ranked = freeform.zip(scores).sortedBy { it.second }.map { it.first }
        return rest + ranked
    }

    /**
     * Collapses assistantToolCall/toolResult pairs that repeat an earlier call's exact (name,
     * arguments) — keeping only the most recent occurrence — so a stalled stage repeating the same
     * read doesn't show the model N copies of an identical question-and-answer. Content is the
     * `ToolCallRequest` JSON built in SessionOrchestrator; parsed generically here (not via the
     * `core:inference` type) to avoid a cross-core dependency. Calls whose arguments differ (e.g.
     * different start_line/end_line, a different path) fingerprint differently and are untouched —
     * that's genuine incremental progress, not repetition.
     */
    private fun dedupeRepeatedToolCalls(entries: List<ContextEntry>): List<ContextEntry> {
        val calls = entries.filter { it.sourceType == "assistantToolCall" }
        val callsById = calls.associateBy { it.sourceId }
        val fingerprints = calls.mapNotNull { entry ->
            toolCallFingerprint(entry.content)?.let { entry.sourceId to it }
        }
        val dropIds = fingerprints.groupBy({ it.second }, { it.first })
            .values
            .filter { it.size > 1 }
            .flatMapTo(mutableSetOf()) { sourceIds ->
                val keepId = sourceIds.maxWith(compareBy(
                    { id: String -> if (callsById.getValue(id).reasoning != null) 1 else 0 },
                    { id: String -> callsById.getValue(id).ordinal }
                ))
                sourceIds.filterNot { it == keepId }
            }
        if (dropIds.isEmpty()) return entries
        return entries.filterNot {
            (it.sourceType == "assistantToolCall" || it.sourceType == "toolResult") && it.sourceId in dropIds
        }
    }

    // A successful tool result is framed "[<tool> exit=<code>]\n..." by renderToolResult; failures
    // use ERROR:/FATAL: sentinels (never matched here). Only exit=0 writes are compacted.
    private val successFrame = Regex("""^\[(\S+) exit=(\d+)]""")
    private val writeToolNames = setOf("file_write", "file_edit")

    private fun compactWriteReceipts(entries: List<ContextEntry>): List<ContextEntry> {
        val writeCallPaths = entries
            .filter { it.sourceType == "assistantToolCall" && toolCallName(it.content) in writeToolNames }
            .associate { it.sourceId to toolCallPath(it.content) }
        if (writeCallPaths.isEmpty()) return entries
        return entries.map { entry ->
            if (entry.sourceType != "toolResult" || entry.sourceId !in writeCallPaths) return@map entry
            val header = entry.content.substringBefore('\n')
            val match = successFrame.find(header) ?: return@map entry
            // A nonzero-exit write is a Success carrying a real advisory (e.g. a partial patch) —
            // keep it verbatim; only a clean exit=0 write body is pure echo we can drop.
            if (match.groupValues[2] != "0") return@map entry
            val path = writeCallPaths[entry.sourceId].orEmpty()
            reencode(entry, "$header wrote $path — succeeded; body elided (full result in event log)".trim())
        }
    }

    private fun toolCallName(content: String): String? = runCatching {
        Json.parseToJsonElement(content).jsonObject["function"]?.jsonObject?.get("name")?.jsonPrimitive?.content
    }.getOrNull()

    private fun toolCallPath(content: String): String? = runCatching {
        val args = Json.parseToJsonElement(content).jsonObject["function"]?.jsonObject
            ?.get("arguments")?.jsonPrimitive?.content ?: return null
        Json.parseToJsonElement(args).jsonObject["path"]?.jsonPrimitive?.content
    }.getOrNull()

    private fun toolCallFingerprint(content: String): String? = runCatching {
        val obj = Json.parseToJsonElement(content).jsonObject
        val fn = obj["function"]?.jsonObject ?: return null
        val name = fn["name"]?.jsonPrimitive?.content ?: return null
        val arguments = fn["arguments"]?.jsonPrimitive?.content ?: ""
        "$name:$arguments"
    }.getOrNull()

    /**
     * A completed tool call is represented as two independently-IDed, independently-compressed
     * entries — assistantToolCall + toolResult — linked only by [ContextEntry.sourceId]. Because
     * `compressed` groups and trims by sourceType (see the `grouped` fold above), one side of a
     * pair can survive its group's budget while the other is trimmed away. An orphaned
     * assistantToolCall with no matching result renders as "I called this tool" with nothing after
     * it — indistinguishable, to the model, from the call silently failing — which drove a live
     * read-loop repro (2026-07-08: discovery stage re-issued the same file_read repeatedly once its
     * result vanished from two consecutive rounds). Drop whichever side lost its pair so every
     * tool-call entry that survives has its result, and vice versa.
     */
    private fun reconcileToolPairs(entries: List<ContextEntry>): List<ContextEntry> {
        val callIds = entries.filter { it.sourceType == "assistantToolCall" }.mapTo(mutableSetOf()) { it.sourceId }
        val resultIds = entries.filter { it.sourceType == "toolResult" }.mapTo(mutableSetOf()) { it.sourceId }
        return entries.filterNot {
            (it.sourceType == "assistantToolCall" && it.sourceId !in resultIds) ||
                (it.sourceType == "toolResult" && it.sourceId !in callIds)
        }
    }

    private fun reencode(entry: ContextEntry, content: String): ContextEntry =
        if (content == entry.content) entry else entry.copy(content = content, tokenEstimate = estimateTokens(content))

    private fun estimateTokens(text: String): Int = (text.length + (CHARS_PER_TOKEN - 1)) / CHARS_PER_TOKEN

    private fun factSheetEntry(content: String) = ContextEntry(
        id = ContextEntryId("factsheet"),
        layer = ContextLayer.L0,
        content = content,
        sourceType = "factSheet",
        sourceId = "factSheet",
        tokenEstimate = estimateTokens(content),
        // #312: re-extracted from the live entry set on EVERY build — the most mutable entry in the
        // pack, so it must not sit in the cached system prefix. USER, pinned at L0 with the lowest
        // ordinal so it still renders ahead of the transcript.
        role = EntryRole.USER,
        ordinal = FACT_SHEET_ORDINAL,
    )

    private fun strategyFor(sourceType: String): CompressionStrategy = when (sourceType) {
        "toolLog" -> CompressionStrategy.ToolLog
        "artifact" -> CompressionStrategy.Artifact
        "steeringNote" -> CompressionStrategy.SteeringNote
        "eventHistory" -> CompressionStrategy.EventHistory
        "decisionJournal" -> CompressionStrategy.EventHistory
        else -> CompressionStrategy.Conversation()
    }
}
