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,
)

object PromptRenderer {
    // 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 (L0 additionally folds regardless of role). 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.
        val (systemEntries, conversationEntries) = contextPack.layers.entries
            .flatMap { (layer, entries) -> entries.map { layer to it } }
            .partition { (layer, entry) -> layer == ContextLayer.L0 || entry.role == EntryRole.SYSTEM }
        val systemContent = systemEntries
            .sortedWith(compareBy({ it.first.ordinal }, { it.second.ordinal }))
            .joinToString("\n\n") { it.second.content }
            .takeIf { it.isNotBlank() }
        val conversationMessages = conversationEntries
            .sortedWith(compareBy({ it.second.ordinal }, { layerPriority(it.first) }))
            .map { (_, entry) -> entry.toChatMessage() }
        // Repetition anchoring: steering directives fold into the leading system message, 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.
        val anchor = systemEntries
            .filter { it.second.sourceType == "steeringNote" }
            .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")) }
        }
        return messages.ifEmpty { listOf(ChatMessage("user", "")) }
    }

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