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

import com.correx.core.context.builder.DefaultContextPackBuilder
import com.correx.core.context.builder.DefaultDecisionPointBuilder
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 com.correx.core.events.types.ContextPackId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test

class DefaultContextPackBuilderTest {

    private val compressor = DefaultContextCompressor()
    private val builder = DefaultContextPackBuilder(compressor)
    private val decisionBuilder = DefaultDecisionPointBuilder()

    private val sessionId = SessionId("s1")
    private val stageId = StageId("stage-1")
    private val packId = ContextPackId("pack-1")

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

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

    @Test
    fun `oversized tool results from a stage tool loop are trimmed to budget`() = kotlinx.coroutines.runBlocking {
        // Live repro (2026-06-11): analyst stage with three file_read results of ~5.6k/11.9k/10.7k
        // tokens sailed through a 16384 budget untrimmed and produced a 34k-token prompt.
        val entries = listOf(
            typedEntry("sys", ContextLayer.L0, 3500, "systemPrompt"),
            typedEntry("task", ContextLayer.L1, 150, "agentPrompt"),
            typedEntry("call1", ContextLayer.L2, 40, "assistantToolCall"),
            typedEntry("res1", ContextLayer.L2, 5600, "toolResult"),
            typedEntry("call2", ContextLayer.L2, 40, "assistantToolCall"),
            typedEntry("res2", ContextLayer.L2, 11900, "toolResult"),
            typedEntry("call3", ContextLayer.L2, 40, "assistantToolCall"),
            typedEntry("res3", ContextLayer.L2, 10700, "toolResult"),
        )
        val pack = builder.build(packId, sessionId, stageId, entries, TokenBudget(limit = 16384))
        assertTrue(
            pack.budgetUsed <= 16384,
            "budgetUsed ${pack.budgetUsed} must not exceed the 16384 stage budget",
        )
    }

    @Test
    fun `build groups entries by layer`() = kotlinx.coroutines.runBlocking {
        val entries = listOf(
            entry("a", ContextLayer.L0, 50),
            entry("b", ContextLayer.L1, 60),
            entry("c", ContextLayer.L2, 40)
        )
        val pack = builder.build(packId, sessionId, stageId, entries, TokenBudget(limit = 4000))
        assertEquals(1, pack.layers[ContextLayer.L0]?.size)
        assertEquals(1, pack.layers[ContextLayer.L1]?.size)
        assertEquals(1, pack.layers[ContextLayer.L2]?.size)
    }

    @Test
    fun `budgetUsed reflects total token estimate of retained entries`() = kotlinx.coroutines.runBlocking {
        val entries = listOf(
            entry("a", ContextLayer.L0, 100),
            entry("b", ContextLayer.L1, 200)
        )
        val pack = builder.build(packId, sessionId, stageId, entries, TokenBudget(limit = 4000))
        assertEquals(300, pack.budgetUsed)
    }

    @Test
    fun `pack budgetUsed never exceeds budgetLimit`() = kotlinx.coroutines.runBlocking {
        val entries = (1..20).map { entry("e$it", ContextLayer.L2, 100) }
        val pack = builder.build(packId, sessionId, stageId, entries, TokenBudget(limit = 500))
        assertTrue(pack.budgetUsed <= 500)
    }

    @Test
    fun `DecisionPointBuilder returns only L0 and L1 entries`() = kotlinx.coroutines.runBlocking {
        val entries = listOf(
            entry("a", ContextLayer.L0, 50),
            entry("b", ContextLayer.L1, 60),
            entry("c", ContextLayer.L2, 40)
        )
        val pack = builder.build(packId, sessionId, stageId, entries, TokenBudget(limit = 4000))
        val decisionEntries = decisionBuilder.build(pack)
        assertTrue(decisionEntries.all { it.layer == ContextLayer.L0 || it.layer == ContextLayer.L1 })
        assertEquals(2, decisionEntries.size)
    }

    @Test
    fun `build handles empty entries list`() = kotlinx.coroutines.runBlocking {
        val pack = builder.build(packId, sessionId, stageId, emptyList(), TokenBudget(limit = 1000))
        assertEquals(0, pack.budgetUsed)
        assertEquals(0, pack.compressionMetadata.entriesDropped)
        assertTrue(pack.layers.isEmpty())
    }

    @Test
    fun `steeringNote and eventHistory entries are retained even when exceeding budget`() = kotlinx.coroutines.runBlocking {
        val entries = listOf(
            entry("steering", ContextLayer.L0, 500).copy(sourceType = "steeringNote"),
            entry("history", ContextLayer.L1, 600).copy(sourceType = "eventHistory"),
            entry("normal", ContextLayer.L2, 200)
        )
        val pack = builder.build(packId, sessionId, stageId, entries, TokenBudget(limit = 800))
        val allRetained = pack.layers.values.flatten()
        assertTrue(allRetained.any { it.id == ContextEntryId("steering") }, "steeringNote must be retained")
        assertTrue(allRetained.any { it.id == ContextEntryId("history") }, "eventHistory must be retained")
    }

    @Test
    fun `tool dialogue keeps chronological order despite sourceType grouping`() = kotlinx.coroutines.runBlocking {
        // A multi-round tool loop: assistant call then tool result, interleaved. Compression
        // groups by sourceType internally; the builder must restore input order so the L2
        // layer reads call, result, call, result — not all calls then all results.
        val entries = listOf(
            entry("task", ContextLayer.L1, 10).copy(sourceType = "agentPrompt"),
            entry("call1", ContextLayer.L2, 10).copy(sourceType = "assistantToolCall"),
            entry("result1", ContextLayer.L2, 10).copy(sourceType = "toolResult"),
            entry("call2", ContextLayer.L2, 10).copy(sourceType = "assistantToolCall"),
            entry("result2", ContextLayer.L2, 10).copy(sourceType = "toolResult"),
        )
        val pack = builder.build(packId, sessionId, stageId, entries, TokenBudget(limit = 4000))
        val l2Ids = pack.layers[ContextLayer.L2]!!.map { it.id.value }
        assertEquals(listOf("call1", "result1", "call2", "result2"), l2Ids)
        // ordinals are stamped from input order
        assertEquals(listOf(1, 2, 3, 4), pack.layers[ContextLayer.L2]!!.map { it.ordinal })
    }
}