correx/testing/deterministic/src/test/kotlin/ContextCompressionDeterminismTest.kt

import com.correx.core.context.compression.CompressionStrategy
import com.correx.core.context.compression.DefaultContextCompressor
import com.correx.core.context.model.ContextEntry
import com.correx.core.context.model.ContextLayer
import com.correx.core.context.model.TokenBudget
import com.correx.core.events.types.ContextEntryId
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test
import org.junit.jupiter.params.ParameterizedTest
import org.junit.jupiter.params.provider.MethodSource

class ContextCompressionDeterminismTest {

    private val compressor = DefaultContextCompressor()

    private fun entry(
        id: String,
        layer: ContextLayer,
        tokens: Int,
        content: String = "content-$id",
        sourceId: String = id
    ) = ContextEntry(
        id = ContextEntryId(id),
        layer = layer,
        content = content,
        sourceType = "test",
            sourceId = sourceId,
            tokenEstimate = tokens
        )

    @Test
    fun `same input always produces same output`() {
        val entries = listOf(
            entry("a", ContextLayer.L2, 100),
            entry("b", ContextLayer.L1, 50),
            entry("c", ContextLayer.L0, 30)
        )
        val budget = TokenBudget(limit = 200)
        val result1 = compressor.compress(entries, budget, CompressionStrategy.Conversation())
        val result2 = compressor.compress(entries, budget, CompressionStrategy.Conversation())
        assertEquals(result1, result2)
    }

    @Test
    fun `L2 entries are dropped before L1 when over budget`() {
        val entries = listOf(
            entry("l2-1", ContextLayer.L2, 80),
            entry("l1-1", ContextLayer.L1, 80),
            entry("l0-1", ContextLayer.L0, 80)
        )
        val budget = TokenBudget(limit = 170)
        val result = compressor.compress(entries, budget, CompressionStrategy.Conversation())
        assertTrue(result.none { it.layer == ContextLayer.L2 }, "L2 must be dropped first")
        assertTrue(result.any { it.layer == ContextLayer.L1 }, "L1 must be retained")
        assertTrue(result.any { it.layer == ContextLayer.L0 }, "L0 must always be retained")
    }

    @Test
    fun `SteeringNote entries are never dropped regardless of budget`() {
        val entries = listOf(
            entry("s1", ContextLayer.L1, 500),
            entry("s2", ContextLayer.L2, 500)
        )
        val budget = TokenBudget(limit = 10)
        val result = compressor.compress(entries, budget, CompressionStrategy.SteeringNote)
        assertEquals(2, result.size)
    }

    @Test
    fun `EventHistory entries are never dropped`() {
        val entries = listOf(
            entry("e1", ContextLayer.L2, 500),
            entry("e2", ContextLayer.L2, 500)
        )
        val budget = TokenBudget(limit = 10)
        val result = compressor.compress(entries, budget, CompressionStrategy.EventHistory)
        assertEquals(2, result.size)
    }

    @Test
    fun `Artifact strategy keeps latest entry per sourceId`() {
        val entries = listOf(
            entry("v1", ContextLayer.L1, 50, sourceId = "artifact-a"),
            entry("v2", ContextLayer.L1, 50, sourceId = "artifact-a"),
            entry("v3", ContextLayer.L1, 50, sourceId = "artifact-b")
        )
        val budget = TokenBudget(limit = 4000)
        val result = compressor.compress(entries, budget, CompressionStrategy.Artifact)
        assertEquals(2, result.size)
        assertTrue(result.any { it.id == ContextEntryId("v2") }, "Latest artifact-a version must be kept")
        assertTrue(result.any { it.id == ContextEntryId("v3") }, "artifact-b must be kept")
    }

    @Test
    fun `ToolLog strategy deduplicates identical content`() {
        val entries = listOf(
            entry("t1", ContextLayer.L1, 30, content = "same output"),
            entry("t2", ContextLayer.L1, 30, content = "same output"),
            entry("t3", ContextLayer.L1, 30, content = "different output")
        )
        val budget = TokenBudget(limit = 4000)
        val result = compressor.compress(entries, budget, CompressionStrategy.ToolLog)
        assertEquals(2, result.size)
    }

    @Test
    fun `entries fitting within budget are not dropped`() {
        val entries = listOf(
            entry("a", ContextLayer.L1, 30),
            entry("b", ContextLayer.L2, 30)
        )
        val budget = TokenBudget(limit = 200)
        val result = compressor.compress(entries, budget, CompressionStrategy.Conversation())
        assertEquals(2, result.size)
    }

    @Test
    fun `Conversation keepLast drops oldest entries beyond the limit`() {
        val entries = (1..5).map { entry("e$it", ContextLayer.L1, 10) }
        val budget = TokenBudget(limit = 4000)
        val result = compressor.compress(entries, budget, CompressionStrategy.Conversation(keepLast = 3))
        assertEquals(3, result.size)
        assertEquals(listOf("e3", "e4", "e5"), result.map { it.id.value })
    }

    @Test
    fun `Conversation L0 is never truncated when L1 overflows budget`() {
        val entries = listOf(
            entry("l0-1", ContextLayer.L0, 60),
            entry("l1-1", ContextLayer.L1, 60),
            entry("l1-2", ContextLayer.L1, 60)
        )
        val budget = TokenBudget(limit = 100)
        val result = compressor.compress(entries, budget, CompressionStrategy.Conversation())
        assertTrue(result.any { it.layer == ContextLayer.L0 }, "L0 must never be dropped")
        assertTrue(result.none { it.id == ContextEntryId("l1-1") }, "Oldest L1 must be dropped first")
    }

    // Parameterized equivalence test: verifies the refactored running-total trimToFit
    // produces the same output as the reference naive implementation for varied input shapes.
    @ParameterizedTest
    @MethodSource("trimToFitEquivalenceCases")
    fun `trimToFit running-total matches reference naive implementation`(
        entries: List<ContextEntry>,
        budget: TokenBudget
    ) {
        val expected = naiveTrimToFit(entries, budget)
        val actual = compressor.compress(entries, budget, CompressionStrategy.Conversation())
        assertEquals(expected, actual, "running-total trimToFit must match naive O(n²) reference")
    }

    // Reference O(n²) implementation kept as spec baseline — intentionally not optimized.
    private fun naiveTrimToFit(entries: List<ContextEntry>, budget: TokenBudget): List<ContextEntry> {
        val dropOrder = listOf(ContextLayer.L2, ContextLayer.L1)
        if (entries.sumOf { it.tokenEstimate } <= budget.limit) return entries
        val mutable = entries.toMutableList()
        for (layer in dropOrder) {
            if (mutable.sumOf { it.tokenEstimate } <= budget.limit) break
            val layerEntries = mutable.filter { it.layer == layer }
            for (e in layerEntries) {
                if (mutable.sumOf { it.tokenEstimate } <= budget.limit) break
                mutable.remove(e)
            }
        }
        return mutable
    }

    companion object {
        @JvmStatic
        fun trimToFitEquivalenceCases(): List<Array<Any>> {
            fun e(id: String, layer: ContextLayer, tokens: Int) =
                ContextEntry(
                    id = ContextEntryId(id),
                    layer = layer,
                    content = "c-$id",
                    sourceType = "test",
                    sourceId = id,
                    tokenEstimate = tokens
                )

            return listOf(
                // All fit — nothing dropped
                arrayOf(
                    listOf(e("a", ContextLayer.L1, 10), e("b", ContextLayer.L2, 10)),
                    TokenBudget(limit = 100)
                ),
                // L2 fully dropped, L1 retained
                arrayOf(
                    listOf(e("x1", ContextLayer.L2, 80), e("x2", ContextLayer.L1, 80), e("x3", ContextLayer.L0, 40)),
                    TokenBudget(limit = 130)
                ),
                // Both L2 and some L1 dropped
                arrayOf(
                    (1..5).map { e("l2-$it", ContextLayer.L2, 30) } +
                        (1..5).map { e("l1-$it", ContextLayer.L1, 30) } +
                        listOf(e("l0-1", ContextLayer.L0, 20)),
                    TokenBudget(limit = 80)
                ),
                // Empty list
                arrayOf(emptyList<ContextEntry>(), TokenBudget(limit = 50)),
                // Single entry exactly at budget
                arrayOf(listOf(e("only", ContextLayer.L2, 50)), TokenBudget(limit = 50))
            )
        }
    }
}