import com.correx.core.context.compression.CompressionPolicy
import com.correx.core.context.compression.CompressionStage
import com.correx.core.context.compression.ContextClass
import com.correx.core.context.compression.ContextClassifier
import com.correx.core.context.compression.FactSheet
import com.correx.core.context.compression.FormatCompressor
import com.correx.core.context.compression.ProtectedSpanTagger
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.types.ContextEntryId
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test

class CompressionPipelineStagesTest {

    private fun entry(sourceType: String, layer: ContextLayer, role: EntryRole, content: String = "x") =
        ContextEntry(
            id = ContextEntryId("e"),
            layer = layer,
            content = content,
            sourceType = sourceType,
            sourceId = "s",
            tokenEstimate = 1,
            role = role,
        )

    @Test
    fun `levels are additive`() {
        val l1 = CompressionPolicy(1)
        assertTrue(l1.enabled(CompressionStage.FORMAT_COMPRESS))
        assertFalse(l1.enabled(CompressionStage.CACHE_STATIC))

        val l7 = CompressionPolicy(7)
        // includes everything from 1..7
        assertTrue(l7.enabled(CompressionStage.FORMAT_COMPRESS))
        assertTrue(l7.enabled(CompressionStage.FACT_SHEET))
        assertTrue(l7.enabled(CompressionStage.PROTECTED_SPAN_TAG))
        assertFalse(l7.enabled(CompressionStage.TOME_MERGE)) // level 8
    }

    @Test
    fun `protected spans capture ids numbers paths and code`() {
        val tagger = ProtectedSpanTagger(extraEntities = listOf("GatedDeltaNet"))
        val spans = tagger.protectedSpans(
            "connect to 10.0.0.1 at path /etc/nginx.conf value 0xFF port = 8080 " +
                "`inline` on GatedDeltaNet"
        )
        assertTrue(spans.any { it == "10.0.0.1" }, spans.toString())
        assertTrue(spans.any { it.contains("/etc/nginx.conf") }, spans.toString())
        assertTrue(spans.any { it == "0xFF" }, spans.toString())
        assertTrue(spans.any { it == "`inline`" }, spans.toString())
        assertTrue(spans.contains("GatedDeltaNet"), spans.toString())
        assertTrue(spans.any { it == "8080" }, spans.toString())
    }

    @Test
    fun `format compressor flattens json and preserves exact values`() {
        val out = FormatCompressor().compress("""{"a":{"b":1},"c":"10.0.0.1"}""")
        assertEquals("a.b = 1\nc = 10.0.0.1", out)
    }

    @Test
    fun `format compressor leaves non-json content untouched`() {
        val out = FormatCompressor().compress("log line\nlog line\nother")
        assertEquals("log line\nlog line\nother", out)
    }

    @Test
    fun `fact sheet extracts newest-first and caps size`() {
        val fs = FactSheet(ProtectedSpanTagger(), maxFacts = 2)
        val entries = listOf(
            entry("chat", ContextLayer.L2, EntryRole.USER, "old 1.1.1.1").copy(ordinal = 0),
            entry("chat", ContextLayer.L2, EntryRole.USER, "new 2.2.2.2 and 3.3.3.3").copy(ordinal = 1),
        )
        val facts = fs.extract(entries)
        // newest entry's spans come first; cap of 2 evicts the oldest entry's IP
        assertEquals(listOf("2.2.2.2", "3.3.3.3"), facts)
    }

    @Test
    fun `builder pins a fact sheet at level 7 and omits it below`() = kotlinx.coroutines.runBlocking {
        val entries = listOf(entry("chat", ContextLayer.L1, EntryRole.USER, "deploy to 10.0.0.1"))
        val compressor = com.correx.core.context.compression.DefaultContextCompressor()

        suspend fun buildAt(level: Int) = com.correx.core.context.builder.DefaultContextPackBuilder(
            compressor, CompressionPolicy(level)
        ).build(
            com.correx.core.events.types.ContextPackId("p"),
            com.correx.core.events.types.SessionId("s"),
            com.correx.core.events.types.StageId("st"),
            entries,
            com.correx.core.context.model.TokenBudget(limit = 4000),
        )

        val atSeven = buildAt(7).layers.values.flatten()
        assertTrue(atSeven.any { it.sourceType == "factSheet" && it.content.contains("10.0.0.1") })
        assertFalse(buildAt(6).layers.values.flatten().any { it.sourceType == "factSheet" })
        Unit
    }

    @Test
    fun `tome merges near-duplicate freeform turns keeping the newest`() {
        val merger = com.correx.core.context.compression.ToMeMerger(threshold = 0.8)
        val a = entry("chat", ContextLayer.L2, EntryRole.USER, "the deploy failed on host web").copy(ordinal = 0)
        val b = entry("chat", ContextLayer.L2, EntryRole.USER, "the deploy failed on host web again").copy(ordinal = 1)
        val c = entry("chat", ContextLayer.L2, EntryRole.USER, "totally unrelated content here").copy(ordinal = 2)
        val out = merger.merge(listOf(a, b, c))
        // a and b are near-duplicates -> only the newest (b) survives; c is distinct
        assertEquals(listOf(1, 2), out.map { it.ordinal })
    }

    @Test
    fun `token pruning at level 3 rewrites freeform preserving protected spans`() = kotlinx.coroutines.runBlocking {
        val pruner = object : com.correx.core.context.compression.TokenPruner {
            override suspend fun prune(content: String, protectedSpans: List<String>, targetRatio: Double): String =
                "PRUNED " + protectedSpans.joinToString(" ")
        }
        val builder = com.correx.core.context.builder.DefaultContextPackBuilder(
            com.correx.core.context.compression.DefaultContextCompressor(),
            CompressionPolicy(3),
            tokenPruner = pruner,
        )
        val entries = listOf(entry("chat", ContextLayer.L1, EntryRole.USER, "please redeploy 10.0.0.1 now"))
        val pack = builder.build(
            com.correx.core.events.types.ContextPackId("p"),
            com.correx.core.events.types.SessionId("s"),
            com.correx.core.events.types.StageId("st"),
            entries,
            com.correx.core.context.model.TokenBudget(limit = 4000),
        )
        val content = pack.layers.values.flatten().first { it.sourceType == "chat" }.content
        assertTrue(content.startsWith("PRUNED"), content)
        assertTrue(content.contains("10.0.0.1"), content)
        Unit
    }

    @Test
    fun `static instruction docs are never token-pruned`() = kotlinx.coroutines.runBlocking {
        // Instruction-doc pruning is retired: curated/procedural docs (project profile, AGENTS.md
        // how-to) must reach the model verbatim — token-pruning fused them into unparseable soup.
        val pruner = object : com.correx.core.context.compression.TokenPruner {
            override suspend fun prune(
                content: String,
                protectedSpans: List<String>,
                targetRatio: Double,
            ): String = "DOC-PRUNED"
        }
        val builder = com.correx.core.context.builder.DefaultContextPackBuilder(
            com.correx.core.context.compression.DefaultContextCompressor(), CompressionPolicy(3), tokenPruner = pruner,
        )
        val bigProfile = "# Project profile\n" + "convention line ".repeat(120) // > old DOC_MIN_CHARS
        val bigInstructions = "# AGENTS.md\n" + "5 stage_complete when done. ".repeat(80)
        val entries = listOf(
            entry("proj", ContextLayer.L0, EntryRole.SYSTEM, bigProfile).copy(sourceType = "projectProfile"),
            entry("sys", ContextLayer.L0, EntryRole.SYSTEM, "you are an agent").copy(sourceType = "systemPrompt"),
            entry("agi", ContextLayer.L0, EntryRole.SYSTEM, bigInstructions).copy(sourceType = "agentInstructions"),
        )
        val flat = builder.build(
            com.correx.core.events.types.ContextPackId("p"),
            com.correx.core.events.types.SessionId("s"),
            com.correx.core.events.types.StageId("st"),
            entries,
            com.correx.core.context.model.TokenBudget(limit = 4000),
        ).layers.values.flatten()
        assertEquals(bigProfile, flat.first { it.sourceType == "projectProfile" }.content)
        assertEquals("you are an agent", flat.first { it.sourceType == "systemPrompt" }.content)
        assertEquals(bigInstructions, flat.first { it.sourceType == "agentInstructions" }.content)
        Unit
    }

    @Test
    fun `classifier separates static structured freeform`() {
        val c = ContextClassifier()
        assertEquals(ContextClass.STATIC, c.classify(entry("systemPrompt", ContextLayer.L0, EntryRole.SYSTEM)))
        assertEquals(ContextClass.STRUCTURED, c.classify(entry("toolLog", ContextLayer.L2, EntryRole.TOOL)))
        assertEquals(ContextClass.STRUCTURED, c.classify(entry("steeringNote", ContextLayer.L2, EntryRole.SYSTEM)))
        // Assistant tool-call turns are structured history, never token-pruned (amnesia-loop fix).
        assertEquals(
            ContextClass.STRUCTURED,
            c.classify(entry("assistantToolCall", ContextLayer.L2, EntryRole.ASSISTANT)),
        )
        assertEquals(ContextClass.FREEFORM, c.classify(entry("chat", ContextLayer.L1, EntryRole.USER)))
    }
}
