package com.correx.core.inference

import com.correx.core.context.model.ContextEntry
import com.correx.core.context.model.ContextLayer
import com.correx.core.context.model.ContextPack
import com.correx.core.context.model.EntryRole
import kotlinx.serialization.Serializable

@Serializable
data class ChatMessage(
    val role: String,
    val content: String,
    // Carries ContextEntry.reasoning forward for assistant tool-call turns — see that field's
    // doc for why this exists (Gemma-family looping when its own prior reasoning is missing).
    val reasoning: String? = null,
    // Optional tool-call-id that links a role:tool response to the assistant's tool_calls entry
    // (sourceId in the event store). Required by the Gemma4 Jinja template to render tool
    // response blocks with the correct function name; without it every tool response is
    // rendered as "unknown" and the thought/reasoning channel logic can break.
    val toolCallId: String? = null,
)

object PromptRenderer {
    // #293: gate retry/recovery repair mandates. Root cause: they used to render as L1/SYSTEM, so
    // they folded into the leading system block — far from the assistant/tool transcript and weaker
    // than the original stage task. Instead they render as the FINAL user message, right after the
    // tool evidence, where a weak local model attends strongest and reads it as the next action.
    // Add a sourceType here (and set the entry's role to USER) to route it to that trailing slot.
    //
    // #312: highest precedence FIRST — at most ONE mandate renders per turn. The trailing slot works
    // because it is scarce and authoritative; a recovery stage on a retry with an unmet delta would
    // otherwise stack three competing "do this next" blocks and the channel becomes noise again.
    private val repairMandatePrecedence = listOf(
        "recoveryTicket",
        "retryFeedback",
        "groundingFeedback",
        "rejectionFeedback",
    )

    // The remaining-delta checklist is not a competing mandate — it is the stage's completion
    // signal ("what is left to make true") — so it appends after whichever mandate won rather
    // than displacing it.
    private const val COMPLETION_SIGNAL = "remainingDelta"

    private val trailingSourceTypes = repairMandatePrecedence.toSet() + COMPLETION_SIGNAL

    // Tiebreak only: when entries carry no chronological ordinal (all 0 — e.g. router
    // chat, which assembles its pack directly), fall back to the old layer priority that
    // renders L1 (the live user turn) last so the template sees a user query at the end.
    // When ordinals are present (orchestrator stage path), they dominate and produce true
    // chronological order — a tool loop reads task → assistant → tool → assistant → …
    private val layerPriority: (ContextLayer) -> Int = { layer ->
        if (layer == ContextLayer.L1) Int.MAX_VALUE else layer.ordinal
    }

    fun render(contextPack: ContextPack): List<ChatMessage> {
        // Every SYSTEM-role entry folds into the single leading system message, whatever its layer.
        // Strict chat templates (e.g. Qwen) reject any system message that is not the first message,
        // so recalled memory, stage summaries, and retrieval entries must never render as standalone
        // system turns.
        //
        // #312: ROLE alone decides the message type — the old `layer == L0 ||` clause is gone. The
        // system block is for content that does not change during a run (prompts, guidance, profiles);
        // anything the run mutates (steering, gate verdicts, tickets, baselines, claimed task) is a
        // user message, both because a mutating system prefix defeats prompt caching and because
        // models under-weight system-folded content relative to the trailing user turn. Layer keeps
        // its own job — budget tier and pin/prune eligibility (see ContextClassifier).
        val (systemEntries, conversationEntries) = contextPack.layers.entries
            .flatMap { (layer, entries) -> entries.map { layer to it } }
            .partition { (_, entry) -> entry.role == EntryRole.SYSTEM }
        val systemContent = systemEntries
            .sortedWith(compareBy({ it.first.ordinal }, { it.second.ordinal }))
            .joinToString("\n\n") { it.second.content }
            .takeIf { it.isNotBlank() }
        // #293: pull repair mandates out of the inline flow — they render once, as the last turn.
        val (repairPairs, inlinePairs) = conversationEntries
            .partition { it.second.sourceType in trailingSourceTypes }
        val conversationMessages = inlinePairs
            .sortedWith(compareBy({ it.second.ordinal }, { layerPriority(it.first) }))
            .map { (_, entry) -> entry.toChatMessage() }
        // Only the highest-precedence mandate present survives; the rest stay out of the prompt
        // entirely (their content is still in the transcript/event log — this slot is not their
        // only carrier). The completion signal appends after it.
        val mandate = repairMandatePrecedence.firstNotNullOfOrNull { sourceType ->
            repairPairs.contentOf(sourceType)
        }
        val repairMandate = listOfNotNull(mandate, repairPairs.contentOf(COMPLETION_SIGNAL))
            .joinToString("\n\n")
            .takeIf { it.isNotBlank() }
        // Repetition anchoring: steering directives render early (they are pinned standing context),
        // far from the final query — weak local models forget them (lost-in-the-middle). Restate
        // them once as a trailing user turn, where models attend strongest. Template-safe: a
        // user message at the end never trips strict system-must-be-first templates. Scans every
        // entry, not just the system fold: since #312 a locked steering note is USER-role too, and
        // both the locked and unlocked paths deserve the same anchor.
        val anchor = (systemEntries + conversationEntries)
            .filter { it.second.sourceType == "steeringNote" }
            .sortedBy { it.second.ordinal }
            .joinToString("\n") { it.second.content }
            .takeIf { it.isNotBlank() }
        val messages = buildList {
            systemContent?.let { add(ChatMessage("system", it)) }
            addAll(conversationMessages)
            anchor?.let { add(ChatMessage("user", "Reminder — active steering directive(s):\n$it")) }
            // The repair mandate is the final message — the model's next action after the transcript.
            repairMandate?.let { add(ChatMessage("user", it)) }
        }
        return messages.ifEmpty { listOf(ChatMessage("user", "")) }
    }

    private fun List<Pair<ContextLayer, ContextEntry>>.contentOf(sourceType: String): String? =
        filter { it.second.sourceType == sourceType }
            .sortedBy { it.second.ordinal }
            .joinToString("\n\n") { it.second.content }
            .takeIf { it.isNotBlank() }

    private fun ContextEntry.toChatMessage(): ChatMessage = ChatMessage(
        role = when (role) {
            EntryRole.SYSTEM -> "system"
            EntryRole.ASSISTANT -> "assistant"
            EntryRole.TOOL -> "tool"
            EntryRole.USER -> "user"
        },
        content = content,
        reasoning = reasoning,
        toolCallId = if (role == EntryRole.TOOL) sourceId else null,
    )
}
