package net.kentonc.si.agent

import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext

/**
 * An ongoing conversation: holds the message history and streams each reply —
 * a port of the Chat struct in src/llm.rs.
 *
 * One `send` may span several requests: if the model asks to call tools, the
 * calls run, their results go back as `tool` messages, and generation continues
 * until it produces a plain answer.
 */
class Agent(
    /** Settings are read at call time so a mid-conversation config edit applies. */
    private val settings: () -> Config,
    private val tools: ToolHost,
    private val confirm: suspend (ToolCall) -> Boolean,
    private val onItem: (Item) -> Unit,
) {
    private val messages = mutableListOf<WireMessage>()

    init {
        reset()
    }

    /** Drop the history and start over from the system prompt. */
    fun reset() {
        messages.clear()
        messages.add(WireMessage.system(SYSTEM_PROMPT))
    }

    /** Number of turns so far, excluding the system prompt. */
    val turnCount: Int get() = messages.size - 1

    /** Append a user message and stream the reply into `onItem`.
     *
     *  If a request fails, every message from this turn is rolled back so the
     *  history stays consistent with what the model actually saw.
     */
    suspend fun send(userMessage: String, images: List<AgentImage>) {
        val mark = messages.size
        messages.add(WireMessage.user(userMessage, images))
        try {
            runTurn()
        } catch (e: Exception) {
            messages.subList(mark, messages.size).clear()
            throw e
        }
    }

    private suspend fun runTurn() {
        for (round in 0 until MAX_TOOL_ROUNDS) {
            val (reply, toolCalls) = streamReply()

            if (toolCalls.isEmpty()) {
                messages.add(WireMessage.assistant(reply, emptyList()))
                return
            }

            // Out of rounds: stop calling tools but still ask for an answer, so
            // everything gathered so far gets summarized rather than lost.
            if (round + 1 == MAX_TOOL_ROUNDS) {
                messages.add(WireMessage.assistant(reply, toolCalls))
                messages.add(WireMessage.user("Tool budget exhausted. Answer using only what you have gathered so far.", emptyList()))
                val (finalReply, _) = streamReply()
                messages.add(WireMessage.assistant(finalReply, emptyList()))
                return
            }

            // Record the assistant turn, tool calls included: the server needs
            // them to match the results that follow.
            messages.add(WireMessage.assistant(reply, toolCalls))

            for (call in toolCalls) {
                onItem(Item.ToolStart(call.function.name, call.summary()))
                val result = tools.execute(call, settings().autoApprove, confirm)
                onItem(Item.ToolDone(outcomeSummary(result)))
                messages.add(WireMessage.toolResult(call.id, result))
            }
        }
    }

    private suspend fun streamReply(): LlmClient.Reply = withContext(Dispatchers.IO) {
        val cfg = settings()
        clientFor(cfg).streamReply(messages, onItem)
    }

    private fun clientFor(cfg: Config) = LlmClient(
        baseUrl = cfg.baseUrl,
        apiKey = cfg.apiKey.ifBlank { "none" },
        model = cfg.model,
        toolsSpec = tools.specs,
        maxTokens = cfg.maxTokens,
        temperature = cfg.temperature,
    )

    companion object {
        /** Cap on tool round-trips within one user turn: a guard against a
         *  model that loops without ever answering (config.rs). */
        const val MAX_TOOL_ROUNDS = 50

        /** System prompt prepended to every request, adapted for Android. */
        const val SYSTEM_PROMPT =
            "You are a helpful assistant running inside an Android app. " +
            "You have tools to read and search files, write files, and run shell commands with sh. " +
            "File paths are relative to the app's working directory, app-private storage on the device. " +
            "Use tools when they would help instead of guessing. " +
            "Answers are rendered as Markdown, so use headings, lists, and bold where they help readability, " +
            "and keep responses concise."
    }
}

/** One line describing how a tool call went, for the status display
 *  (outcome_summary in llm.rs). */
fun outcomeSummary(result: String): String {
    val trimmed = result.trimStart()
    if (trimmed.startsWith("Error:")) {
        val rest = trimmed.removePrefix("Error:").trim().take(70)
        return "error: $rest"
    }
    if (trimmed.startsWith("Skipped:")) return "skipped"
    val lines = result.count { it == '\n' } + 1
    val first = result.trim().lineSequence().firstOrNull().orEmpty().take(48)
    return "$lines line(s) — $first"
}