package com.correx.core.context.compression

import com.correx.core.context.model.ContextEntry
import com.correx.core.context.model.ContextLayer
import com.correx.core.context.model.TokenBudget

class DefaultContextCompressor : ContextCompressor {

    // L0 (live execution), L3 (durable project memory), and L4 (archival) are never dropped — only L2 then L1.
    private val dropOrder = listOf(ContextLayer.L2, ContextLayer.L1)

    override fun compress(
        entries: List<ContextEntry>,
        budget: TokenBudget,
        strategy: CompressionStrategy
    ): List<ContextEntry> = when (strategy) {
        CompressionStrategy.SteeringNote,
        CompressionStrategy.EventHistory -> entries
        CompressionStrategy.Artifact -> compressArtifacts(entries, budget)
        CompressionStrategy.ToolLog -> compressToolLogs(entries, budget)
        is CompressionStrategy.Conversation -> compressConversation(entries, budget, strategy.keepLast)
    }

    // Keeps the last `keepLast` entries (oldest-first order preserved), then enforces budget.
    // L0 is never truncated — overflow removes L1 from oldest first.
    private fun compressConversation(
        entries: List<ContextEntry>,
        budget: TokenBudget,
        keepLast: Int
    ): List<ContextEntry> {
        val kept = if (entries.size > keepLast) entries.takeLast(keepLast) else entries
        return trimToFit(kept, budget)
    }

    private fun compressArtifacts(entries: List<ContextEntry>, budget: TokenBudget): List<ContextEntry> {
        val deduplicated = entries
            .groupBy { it.sourceId }
            .map { (_, group) -> group.last() }
        return trimToFit(deduplicated, budget)
    }

    private fun compressToolLogs(entries: List<ContextEntry>, budget: TokenBudget): List<ContextEntry> {
        val deduplicated = entries.distinctBy { it.content }
        return trimToFit(deduplicated, budget)
    }

    private fun trimToFit(entries: List<ContextEntry>, budget: TokenBudget): List<ContextEntry> {
        var total = entries.sumOf { it.tokenEstimate }
        if (total <= budget.limit) return entries

        val mutable = entries.toMutableList()
        for (layer in dropOrder) {
            if (total <= budget.limit) break
            val layerEntries = mutable.filter { it.layer == layer }
            for (entry in layerEntries) {
                if (total <= budget.limit) break
                mutable.remove(entry)
                total -= entry.tokenEstimate
            }
        }
        return mutable
    }
}
