package com.correx.core.talkie

import com.correx.core.context.model.CompressionMetadata
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 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 com.correx.core.inference.Tokenizer
import com.correx.core.talkie.model.Idea
import com.correx.core.talkie.model.NarrationTrigger
import com.correx.core.talkie.model.TalkieConfig
import com.correx.core.talkie.model.RouterL2Entry
import com.correx.core.talkie.model.TalkieState
import com.correx.core.talkie.model.TalkieTurn
import com.correx.core.talkie.model.TurnRole
import com.correx.core.talkie.model.WorkflowSummary
import java.util.*
import kotlinx.coroutines.CancellationException

interface TalkieContextBuilder {
    suspend fun build(
        state: TalkieState,
        budget: TokenBudget,
        availableWorkflows: List<WorkflowSummary> = emptyList(),
        projectProfileText: String? = null,
    ): ContextPack

    /**
     * Builds a minimal [ContextPack] for a narration inference call.
     * Implementations should include system prompt, workflow status, recent L2, and the trigger
     * instruction, but must NOT include conversation history or L3 recalled memory.
     *
     * Default delegates to [build] so that anonymous test stubs that only override [build] compile
     * without needing a stub override.
     */
    suspend fun buildNarrationContext(
        state: TalkieState,
        trigger: NarrationTrigger,
        budget: TokenBudget,
        recentActivity: String? = null,
    ): ContextPack = build(state, budget, emptyList(), null)
}

class DefaultTalkieContextBuilder(
    private val config: TalkieConfig,
    private val tokenizer: Tokenizer? = null,
    // Cross-session ideas to surface to the router (router-only context). Read fresh each build so
    // newly-captured ideas appear on the next turn; injected via the constructor so the build()
    // interface — and its many test doubles — stay unchanged.
    private val ideaProvider: () -> List<Idea> = { emptyList() },
) : TalkieContextBuilder {

    companion object {
        private val SYSTEM_PROMPT =
            """
            You are correx's router — the conversational interface between the operator and a
            local, event-sourced workflow engine. You are not the engine: stages, tools, and
            approvals are run by the kernel and recorded as events. Your job is to help the
            operator understand and steer what the engine is doing.

            Ground every statement in the workflow state and conversation context you are given.
            Never invent stages, tool results, file contents, or status you have not been shown.
            Be concise and direct — plain sentences, no filler. When the operator is steering
            an approval, help them state clearly what they want changed.

            When the operator's request is not something you can answer or steer directly:
            do not simply refuse. Triage instead, citing only what appears in your context:
            (a) if "## Available workflows" is present, suggest the best-fitting workflow(s) and
                say in prose what each would accomplish for this request. Then, on a final line,
                append a fenced json block offering them as a picker:
                ```json
                {"propose_workflows":[{"id":"<workflow-id>","reason":"<one line>"}],"prompt":"<short question>"}
                ```
                Use ONLY ids from "## Available workflows"; list them best-first. The operator
                gets a choice panel (with a manual-answer slot) — omit the block entirely if none fit;
            (b) if "## Project profile" is present and the request relates to the project,
                point at the relevant conventions or commands;
            (c) if recalled cross-session memory shows similar prior work, say so and reference it;
            (d) otherwise offer to think the problem through together in chat.

            Rubber-duck capture: while thinking a problem through in chat, when a concrete idea or
            decision crystallizes, capture it by ending your reply with a fenced json block:
            ```json
            {"ideas":["<one idea per item, in the operator's framing>"]}
            ```
            These are saved to the operator's idea board (also shown back to you under "## Ideas").
            Capture only genuinely new, concrete ideas — omit the block when nothing crystallized,
            and never re-capture an idea already under "## Ideas".
            """.trimIndent()

        private val NARRATION_SYSTEM_PROMPT =
            """
            You are correx's narrator, describing a running workflow to the operator in their live
            feed. In one or two short sentences, in your own voice, say what just happened and
            what comes next. Be concrete and grounded in the state you are given — name the stage,
            the specific actions taken (the files read or written, tools run), the outcome, and the
            reason for any failure or pause. Prefer the concrete actions over restating the stage
            name. Do not use lists, headings, or preamble, and do not invent details you were not
            given. Speak in the present, as events unfold.
            """.trimIndent()

        private const val RECALLED_MEMORY_PREFIX = "[recalled memory]"

        /** Cap on ideas folded into router context, newest-first, so the board can't balloon L0. */
        private const val MAX_IDEAS_IN_CONTEXT = 10

        /** Rough chars-per-token heuristic used when no tokenizer is available. */
        private const val CHARS_PER_TOKEN_ESTIMATE = 4
    }

    override suspend fun build(
        state: TalkieState,
        budget: TokenBudget,
        availableWorkflows: List<WorkflowSummary>,
        projectProfileText: String?,
    ): ContextPack {
        // Protected frame: system prompt and workflow status are ALWAYS included regardless of budget.
        val systemPrompt = buildContextEntry(
            sourceType = "systemPrompt",
            sourceId = "router-system",
            content = SYSTEM_PROMPT,
            layer = ContextLayer.L0,
            role = EntryRole.SYSTEM,
        )
        val workflowStatusEntry = buildContextEntry(
            sourceType = "workflowStatus",
            sourceId = state.currentStageId?.value ?: "none",
            content = buildWorkflowStatusContent(state),
            layer = ContextLayer.L0,
            role = EntryRole.SYSTEM,
        )
        val availableWorkflowsEntry: ContextEntry? = availableWorkflows
            .takeIf { it.isNotEmpty() }
            ?.let {
                buildContextEntry(
                    sourceType = "availableWorkflows",
                    sourceId = "available-workflows",
                    content = buildAvailableWorkflowsContent(it),
                    layer = ContextLayer.L0,
                    role = EntryRole.SYSTEM,
                )
            }
        val projectProfileEntry: ContextEntry? = projectProfileText
            ?.takeIf { it.isNotBlank() }
            ?.let {
                buildContextEntry(
                    sourceType = "projectProfile",
                    sourceId = "project-profile",
                    content = it,
                    layer = ContextLayer.L0,
                    role = EntryRole.SYSTEM,
                )
            }
        val ideasEntry: ContextEntry? = ideaProvider()
            .takeIf { it.isNotEmpty() }
            ?.let {
                buildContextEntry(
                    sourceType = "ideas",
                    sourceId = "operator-ideas",
                    content = buildIdeasContent(it),
                    layer = ContextLayer.L0,
                    role = EntryRole.SYSTEM,
                )
            }

        // Compute remaining budget after protected frame.
        val protectedTokens = systemPrompt.tokenEstimate + workflowStatusEntry.tokenEstimate +
            (availableWorkflowsEntry?.tokenEstimate ?: 0) + (projectProfileEntry?.tokenEstimate ?: 0) +
            (ideasEntry?.tokenEstimate ?: 0)
        var remainingBudget = (budget.limit - protectedTokens).coerceAtLeast(0)

        val allEntries = mutableListOf<ContextEntry>()
        allEntries += systemPrompt
        allEntries += workflowStatusEntry
        availableWorkflowsEntry?.let { allEntries += it }
        projectProfileEntry?.let { allEntries += it }
        ideasEntry?.let { allEntries += it }

        var droppedCount = 0
        val truncatedLayers = mutableSetOf<ContextLayer>()

        // --- Conversation turns (L1) ---
        // The current (protected) user turn is the last element in conversationHistory,
        // regardless of conversationKeepLast — it must always be present.
        val protectedUserTurn: TalkieTurn? = state.conversationHistory.lastOrNull()

        // Apply conversationKeepLast cap. The protected turn may or may not fall within this window.
        val cappedTurns = state.conversationHistory.takeLast(config.conversationKeepLast)

        // Build entries for the capped window, excluding the protected turn (to avoid double-counting).
        val nonProtectedCapped = if (protectedUserTurn != null && cappedTurns.lastOrNull() == protectedUserTurn) {
            cappedTurns.dropLast(1)
        } else {
            cappedTurns
        }

        // Build the protected user turn entry.
        val protectedTurnEntry: ContextEntry? = protectedUserTurn?.let { turn ->
            val role = when (turn.role) {
                TurnRole.USER -> EntryRole.USER
                TurnRole.ROUTER -> EntryRole.ASSISTANT
            }
            buildContextEntry(
                sourceType = "conversation",
                sourceId = "${turn.role.name}-${turn.hashCode()}",
                content = turn.content,
                layer = ContextLayer.L1,
                role = role,
            )
        }

        // Reserve budget for the protected user turn first.
        if (protectedTurnEntry != null) {
            remainingBudget = (remainingBudget - protectedTurnEntry.tokenEstimate).coerceAtLeast(0)
        }

        // Fit remaining capped turns newest-to-oldest (excluding protected), so oldest get dropped when tight.
        val nonProtectedTurnEntries = nonProtectedCapped.map { turn ->
            val role = when (turn.role) {
                TurnRole.USER -> EntryRole.USER
                TurnRole.ROUTER -> EntryRole.ASSISTANT
            }
            buildContextEntry(
                sourceType = "conversation",
                sourceId = "${turn.role.name}-${turn.hashCode()}",
                content = turn.content,
                layer = ContextLayer.L1,
                role = role,
            )
        }

        val fittedTurns = mutableListOf<ContextEntry>()
        for (entry in nonProtectedTurnEntries.asReversed()) {
            if (remainingBudget >= entry.tokenEstimate) {
                remainingBudget -= entry.tokenEstimate
                fittedTurns.add(entry)
            } else {
                droppedCount++
                truncatedLayers.add(ContextLayer.L1)
            }
        }

        // Re-order fitted turns oldest-to-newest (reverse back) and append protected turn at end.
        fittedTurns.reverse()
        allEntries.addAll(fittedTurns)
        protectedTurnEntry?.let { allEntries.add(it) }

        // --- L3 memory (recalled cross-session) ---
        // Sort by score descending (highest relevance first), add what fits.
        val sortedHits = state.lastRetrievedMemory.sortedByDescending { it.score }
        for (hit in sortedHits) {
            val content = "$RECALLED_MEMORY_PREFIX ${hit.text}"
            val entry = buildContextEntry(
                sourceType = "recalledMemory",
                sourceId = hit.entryId,
                content = content,
                layer = ContextLayer.L3,
                role = EntryRole.SYSTEM,
            )
            if (remainingBudget >= entry.tokenEstimate) {
                remainingBudget -= entry.tokenEstimate
                allEntries += entry
            } else {
                droppedCount++
                truncatedLayers.add(ContextLayer.L3)
            }
        }

        // --- L2 memory (compressed session summaries) ---
        // Iterate newest-to-oldest so that when budget is tight the OLDEST entries are dropped.
        val fittedL2 = mutableListOf<ContextEntry>()
        for (l2Entry in state.l2Memory.asReversed()) {
            val content = buildL2Content(l2Entry)
            val entry = buildContextEntry(
                sourceType = "stageSummary",
                sourceId = l2Entry.stageId.value,
                content = content,
                layer = ContextLayer.L2,
                role = EntryRole.SYSTEM,
            )
            if (remainingBudget >= entry.tokenEstimate) {
                remainingBudget -= entry.tokenEstimate
                fittedL2.add(entry)
            } else {
                droppedCount++
                truncatedLayers.add(ContextLayer.L2)
            }
        }
        // Restore chronological (oldest-to-newest) order.
        fittedL2.reverse()
        allEntries.addAll(fittedL2)

        // Build layers map sorted by ContextLayer ordinal for deterministic ordering.
        val layers = allEntries
            .groupBy { it.layer }
            .toSortedMap(compareBy { it.ordinal })

        val budgetUsed = allEntries.sumOf { it.tokenEstimate }

        // Build appliedStrategies list — deterministic, stable order.
        val appliedStrategies = buildList {
            add("L0Immutable")
            add("Conversation")
            if (state.lastRetrievedMemory.isNotEmpty()) add("L3Recall")
            if (budgetUsed > budget.limit) add("BudgetExceeded")
        }

        return ContextPack(
            id = ContextPackId("${state.sessionId?.value ?: "unknown"}-router-pack"),
            sessionId = state.sessionId ?: SessionId("unknown"),
            stageId = state.currentStageId ?: StageId.NONE,
            layers = layers,
            budgetUsed = budgetUsed,
            budgetLimit = budget.limit,
            compressionMetadata = CompressionMetadata(
                appliedStrategies = appliedStrategies,
                truncatedLayers = truncatedLayers.toList(),
                entriesDropped = droppedCount,
            ),
        )
    }

    private fun buildWorkflowStatusContent(state: TalkieState, stageOverride: String? = null): String {
        val status = capitalizeFirst(state.workflowStatus.name.lowercase())
        val stage = stageOverride ?: state.currentStageId?.value ?: "none"
        return "Status: $status, Stage: $stage"
    }

    private fun capitalizeFirst(s: String): String {
        return s.lowercase().let { it[0].uppercaseChar() + it.drop(1) }
    }

    private fun buildL2Content(l2Entry: RouterL2Entry): String {
        val base = "Stage ${l2Entry.stageId.value} (${l2Entry.outcome}): ${l2Entry.summary}"
        val reason = l2Entry.reason?.let { " — reason: $it" }.orEmpty()
        val steering = if (l2Entry.steeringNotes.isNotEmpty()) {
            " | steering: " + l2Entry.steeringNotes.joinToString("; ")
        } else {
            ""
        }
        return base + reason + steering
    }

    private fun buildAvailableWorkflowsContent(workflows: List<WorkflowSummary>): String =
        buildString {
            append("## Available workflows\n")
            workflows.forEach { wf ->
                append("- ${wf.id}: ${wf.description} (stages: ${wf.stageIds.joinToString(" → ")})\n")
            }
        }.trimEnd()

    private fun buildIdeasContent(ideas: List<Idea>): String =
        buildString {
            append("## Ideas\n")
            append("Ideas the operator captured while thinking out loud. Reference them when relevant; ")
            append("they are notes, not instructions.\n")
            ideas.take(MAX_IDEAS_IN_CONTEXT).forEach { append("- ${it.text}\n") }
        }.trimEnd()

    private suspend fun buildContextEntry(
        sourceType: String,
        sourceId: String,
        content: String,
        layer: ContextLayer = ContextLayer.L0,
        role: EntryRole = EntryRole.USER,
    ): ContextEntry {
        return ContextEntry(
            id = ContextEntryId(UUID.randomUUID().toString()),
            layer = layer,
            content = content,
            sourceType = sourceType,
            sourceId = sourceId,
            tokenEstimate = estimateTokens(content),
            role = role,
        )
    }

    private suspend fun estimateTokens(content: String): Int {
        val t = tokenizer
        if (t != null) {
            return runCatching { t.countTokens(content) }.getOrElse { e ->
                if (e is CancellationException) throw e
                fallbackEstimate(content)
            }
        }
        return fallbackEstimate(content)
    }

    private fun fallbackEstimate(content: String): Int {
        return (content.length / CHARS_PER_TOKEN_ESTIMATE).coerceAtLeast(1)
    }

    /**
     * Builds a minimal [ContextPack] for a single narration inference call.
     *
     * Contains (in order, subject to [budget]):
     * 1. L0 — system prompt (always present)
     * 2. L0 — workflow status (always present)
     * 3. L2 — recent stage summaries (oldest dropped first when tight)
     * 4. L0 — trigger instruction (always present as the final user-facing entry)
     *
     * Conversation history and L3 recalled memory are intentionally excluded so
     * narration lines never enter the conversation flow.
     */
    override suspend fun buildNarrationContext(
        state: TalkieState,
        trigger: NarrationTrigger,
        budget: TokenBudget,
        recentActivity: String?,
    ): ContextPack {
        val systemPromptEntry = buildContextEntry(
            sourceType = "systemPrompt",
            sourceId = "router-narration-system",
            content = NARRATION_SYSTEM_PROMPT,
            layer = ContextLayer.L0,
            role = EntryRole.SYSTEM,
        )
        val workflowStatusEntry = buildContextEntry(
            sourceType = "workflowStatus",
            sourceId = trigger.stageId ?: state.currentStageId?.value ?: "none",
            content = buildWorkflowStatusContent(state, trigger.stageId),
            layer = ContextLayer.L0,
            role = EntryRole.SYSTEM,
        )
        // Concrete actions taken (tool calls), folded into the USER instruction below rather than a
        // separate system entry: small narrator models attend to the user turn far more than system
        // context, so this is what actually makes the narration name what happened instead of a
        // generic "stage completed". Null/blank when nothing ran yet.
        val activityLine = recentActivity?.takeIf { it.isNotBlank() }
        // L1, not L0: PromptRenderer folds every L0 entry into the single system message
        // regardless of role, so an L0 user turn would vanish into the system block and the
        // request would go out with no user message (the chat template then rejects it).
        val triggerEntry = buildContextEntry(
            sourceType = "narrationTrigger",
            sourceId = trigger.kind,
            content = activityLine?.let { "${trigger.instruction}\n\n$it" } ?: trigger.instruction,
            layer = ContextLayer.L1,
            role = EntryRole.USER,
        )

        val protectedTokens = systemPromptEntry.tokenEstimate +
            workflowStatusEntry.tokenEstimate +
            triggerEntry.tokenEstimate
        var remainingBudget = (budget.limit - protectedTokens).coerceAtLeast(0)

        val allEntries = mutableListOf<ContextEntry>()
        allEntries += systemPromptEntry
        allEntries += workflowStatusEntry

        var droppedCount = 0
        val truncatedLayers = mutableSetOf<ContextLayer>()

        // L2 stage summaries — newest-to-oldest so oldest get dropped first when tight.
        val fittedL2 = mutableListOf<ContextEntry>()
        for (l2Entry in state.l2Memory.asReversed()) {
            val content = buildL2Content(l2Entry)
            val entry = buildContextEntry(
                sourceType = "stageSummary",
                sourceId = l2Entry.stageId.value,
                content = content,
                layer = ContextLayer.L2,
                role = EntryRole.SYSTEM,
            )
            if (remainingBudget >= entry.tokenEstimate) {
                remainingBudget -= entry.tokenEstimate
                fittedL2.add(entry)
            } else {
                droppedCount++
                truncatedLayers.add(ContextLayer.L2)
            }
        }
        fittedL2.reverse()
        allEntries.addAll(fittedL2)

        // Trigger instruction is always last.
        allEntries += triggerEntry

        val layers = allEntries
            .groupBy { it.layer }
            .toSortedMap(compareBy { it.ordinal })

        val budgetUsed = allEntries.sumOf { it.tokenEstimate }

        return ContextPack(
            id = ContextPackId("${state.sessionId?.value ?: "unknown"}-narration-pack"),
            sessionId = state.sessionId ?: SessionId("unknown"),
            stageId = state.currentStageId ?: StageId.NONE,
            layers = layers,
            budgetUsed = budgetUsed,
            budgetLimit = budget.limit,
            compressionMetadata = CompressionMetadata(
                appliedStrategies = listOf("L0Immutable", "NarrationContext"),
                truncatedLayers = truncatedLayers.toList(),
                entriesDropped = droppedCount,
            ),
        )
    }
}
