package com.correx.core.context.builder

import com.correx.core.context.compression.CompressionStrategy
import com.correx.core.context.compression.ContextCompressor
import com.correx.core.context.model.CompressionMetadata
import com.correx.core.context.model.ContextEntry
import com.correx.core.context.model.ContextPack
import com.correx.core.context.model.TokenBudget
import com.correx.core.events.types.ContextPackId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId

class DefaultContextPackBuilder(
    private val compressor: ContextCompressor
) : ContextPackBuilder {

    private val neverDropSourceTypes = setOf("steeringNote", "eventHistory", "decisionJournal")

    override fun build(
        id: ContextPackId,
        sessionId: SessionId,
        stageId: StageId,
        entries: List<ContextEntry>,
        budget: TokenBudget
    ): ContextPack {
        // Stamp chronological order from the caller's input sequence. Grouping by
        // sourceType (for compression) and by layer reorders entries; the ordinal lets
        // us restore true turn order afterwards so a tool loop reads assistant→tool→…
        // instead of all-assistants-then-all-tools.
        val ordered = entries.mapIndexed { index, entry -> entry.copy(ordinal = index) }
        val (pinned, compressible) = ordered.partition { it.sourceType in neverDropSourceTypes }
        val pinnedTokens = pinned.sumOf { it.tokenEstimate }
        var remainingTokens = (budget.limit - pinnedTokens).coerceAtLeast(0)

        // Dispatch compression strategy by entry sourceType — applying Conversation
        // strategy uniformly mangles tool logs and artifacts (their shape isn't conversational).
        val strategiesUsed = linkedSetOf<String>()
        val compressed = compressible
            .groupBy { it.sourceType }
            .flatMap { (sourceType, group) ->
                val strategy = strategyFor(sourceType)
                strategiesUsed += strategy::class.simpleName ?: "Unknown"
                val groupBudget = TokenBudget(limit = remainingTokens.coerceAtLeast(0))
                val result = compressor.compress(group, groupBudget, strategy)
                remainingTokens -= result.sumOf { it.tokenEstimate }
                result
            }

        // Restore chronological order before layering — groupBy preserves encounter
        // order, so each layer's list comes out in true turn order.
        val retained = (pinned + compressed).sortedBy { it.ordinal }
        val layers = retained.groupBy { it.layer }
        val budgetUsed = retained.sumOf { it.tokenEstimate }
        val droppedCount = entries.size - retained.size
        val truncatedLayers = if (droppedCount > 0) {
            entries.map { it.layer }.distinct().filter { layer ->
                entries.count { it.layer == layer } > retained.count { it.layer == layer }
            }
        } else emptyList()

        return ContextPack(
            id = id,
            sessionId = sessionId,
            stageId = stageId,
            layers = layers,
            budgetUsed = budgetUsed,
            budgetLimit = budget.limit,
            compressionMetadata = CompressionMetadata(
                appliedStrategies = strategiesUsed.toList().ifEmpty { listOf("Conversation") },
                truncatedLayers = truncatedLayers,
                entriesDropped = droppedCount
            )
        )
    }

    private fun strategyFor(sourceType: String): CompressionStrategy = when (sourceType) {
        "toolLog" -> CompressionStrategy.ToolLog
        "artifact" -> CompressionStrategy.Artifact
        "steeringNote" -> CompressionStrategy.SteeringNote
        "eventHistory" -> CompressionStrategy.EventHistory
        "decisionJournal" -> CompressionStrategy.EventHistory
        else -> CompressionStrategy.Conversation()
    }
}
