package com.correx.core.context.compression

import kotlinx.serialization.Serializable

/**
 * The pipeline stages, additively enabled by compression level (see docs/plans/
 * correx-compression-pipeline.md §5). Level N includes every stage of levels 1..N-1.
 */
enum class CompressionStage {
    FORMAT_COMPRESS,            // 1: json→toml, dedupe, flatten. free, deterministic
    CACHE_STATIC,               // 2: kv-cache system prompt / tool defs
    TOKEN_PRUNE,                // 3: LLMLingua-2 on freeform (~45%)
    TIER_SPLIT,                 // 4: age-tiering (recent full, older pruned)
    RECURSIVE_SUMMARIZE,        // 5: tier2/3 summarization
    PROTECTED_SPAN_TAG,         // 6: entity/regex protection before any lossy stage
    FACT_SHEET,                 // 7: standing state block, survives all tiers
    TOME_MERGE,                 // 8: merge near-duplicate turns
    BUDGET_ALLOCATOR_STRICT,    // 9: hard per-tier budget enforcement
}

/**
 * Declarative selection of which pipeline stages run. Additive: raising the level turns
 * on more stages without removing earlier ones. Level 2 (format + static cache) is the
 * always-on default — free and deterministic.
 */
@Serializable
data class CompressionPolicy(val level: Int = DEFAULT_LEVEL) {

    val stages: Set<CompressionStage>
        get() = LEVEL_STAGES.filterKeys { it <= level }.values.flatten().toSet()

    fun enabled(stage: CompressionStage): Boolean = stage in stages

    companion object {
        const val DEFAULT_LEVEL = 2

        // Stage introduced *at* each level; the getter folds in all lower levels.
        private val LEVEL_STAGES: Map<Int, List<CompressionStage>> = mapOf(
            1 to listOf(CompressionStage.FORMAT_COMPRESS),
            2 to listOf(CompressionStage.CACHE_STATIC),
            3 to listOf(CompressionStage.TOKEN_PRUNE),
            4 to listOf(CompressionStage.TIER_SPLIT),
            5 to listOf(CompressionStage.RECURSIVE_SUMMARIZE),
            6 to listOf(CompressionStage.PROTECTED_SPAN_TAG),
            7 to listOf(CompressionStage.FACT_SHEET),
            8 to listOf(CompressionStage.TOME_MERGE),
            9 to listOf(CompressionStage.BUDGET_ALLOCATOR_STRICT),
        )
    }
}
