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> {
        val sorted = contextPack.layers.entries.sortedBy { it.key.ordinal }
        val systemContent = sorted
            .filter { it.key == ContextLayer.L0 }
            .flatMap { it.value }
            .joinToString("\n\n") { it.content }
            .takeIf { it.isNotBlank() }
        val conversationMessages = contextPack.layers.entries
            .filter { it.key != ContextLayer.L0 }
            .flatMap { (layer, entries) -> entries.map { layer to it } }
            .sortedWith(compareBy({ it.second.ordinal }, { layerPriority(it.first) }))
            .map { (_, entry) -> entry.toChatMessage() }
        val messages = buildList {
            systemContent?.let { add(ChatMessage("system", it)) }
            addAll(conversationMessages)
        }
        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,
    )
}
