package net.kentonc.si.agent

import kotlinx.serialization.json.JsonElement
import okhttp3.MediaType.Companion.toMediaType
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.RequestBody.Companion.toRequestBody
import java.io.IOException
import java.io.BufferedReader
import java.io.InputStreamReader
import java.util.TreeMap
import java.util.concurrent.TimeUnit

/**
 * Minimal OpenAI-compatible chat client (llama-server / LiteLLM / vLLM), a port
 * of the streaming half of src/llm.rs. One call posts the whole history and
 * reads the SSE stream; tool-call deltas arrive as fragments spread across many
 * frames and are reassembled here, keyed by their wire index.
 */
class LlmClient(
    private val baseUrl: String,
    private val apiKey: String,
    private val model: String,
    private val toolsSpec: JsonElement,
    private val maxTokens: Int,
    private val temperature: Float,
) {
    /** The instruction to send and everything the caller collects from the stream. */
    data class Reply(val text: String, val toolCalls: List<ToolCall>)

    /** Stream one completion into `onItem`, returning the accumulated reply text
     *  and any fragmented tool calls that arrived complete. */
    fun streamReply(
        messages: List<WireMessage>,
        onItem: (Item) -> Unit,
    ): Reply {
        val body = json.encodeToString(
            ChatRequest.serializer(),
            ChatRequest(
                model = model,
                messages = messages,
                tools = toolsSpec,
                maxTokens = maxTokens,
                temperature = temperature,
                stream = true,
            ),
        )

        val url = "${baseUrl.trimEnd('/')}/v1/chat/completions"
        val request = Request.Builder()
            .url(url)
            .header("Authorization", "Bearer $apiKey")
            .post(body.toRequestBody("application/json".toMediaType()))
            .build()

        val response = client.newCall(request).execute()
        response.use { resp ->
            if (!resp.isSuccessful) {
                val errorText = resp.body?.string()?.take(4000).orEmpty()
                throw IOException("HTTP ${resp.code}: $errorText")
            }
            val source = resp.body?.byteStream()
                ?: throw IOException("empty response body")
            val reader = BufferedReader(InputStreamReader(source, Charsets.UTF_8))

            val reply = StringBuilder()
            val accumulator = ToolCallAccumulator()

            loop@ while (true) {
                val line = reader.readLine() ?: break  // connection ended
                if (!line.startsWith("data:")) continue  // comments, "event:" lines, blank lines
                val data = line.removePrefix("data:").trim()
                if (data == "[DONE]") break@loop

                val chunk = try { json.decodeFromString(StreamChunk.serializer(), data) }
                catch (e: Exception) { null } ?: continue  // tolerate frames we don't understand

                for (choice in chunk.choices) {
                    val delta = choice.delta ?: continue
                    val text = flattenContent(delta.content)
                    if (!text.isNullOrEmpty()) {
                        reply.append(text)
                        onItem(Item.Text(text))
                    }
                    accumulator.push(delta.toolCalls)
                }
            }
            return Reply(reply.toString(), accumulator.finish())
        }
    }

    /** Same shape as the Rust ChatRequest. */
    @kotlinx.serialization.Serializable
    private data class ChatRequest(
        val model: String,
        val messages: List<WireMessage>,
        val tools: JsonElement?,
        @kotlinx.serialization.SerialName("max_tokens") val maxTokens: Int,
        val temperature: Float,
        val stream: Boolean,
    )

    companion object {
        /** One client for the whole app process: connection pool and threads reused. */
        val client: OkHttpClient = OkHttpClient.Builder()
            .connectTimeout(30, TimeUnit.SECONDS)
            // Per-chunk idle timeout, not a total cap: long generations are fine,
            // but a stalled connection fails instead of hanging forever.
            .readTimeout(120, TimeUnit.SECONDS)
            .build()
    }
}

/** Accumulates fragmented tool calls, keyed by their wire index (the Rust
 *  ToolCallAccumulator): the name and id show up once, `arguments` is a JSON
 *  string built up a few characters at a time. */
class ToolCallAccumulator {
    private val calls = TreeMap<Int, ToolCall>()

    fun push(deltas: List<ToolCallDelta>) {
        for (delta in deltas) {
            var call = calls.getOrPut(delta.index) {
                ToolCall(function = FunctionCall(name = "", arguments = ""))
            }
            delta.id?.let { call = call.copy(id = it) }
            delta.function?.let { fn ->
                if (fn.name != null) call = call.copy(function = call.function.copy(name = fn.name))
                if (fn.arguments != null) {
                    call = call.copy(function = call.function.copy(arguments = call.function.arguments + fn.arguments))
                }
            }
            calls[delta.index] = call
        }
    }

    fun finish(): List<ToolCall> = calls.values.mapIndexed { n, call ->
        // Some servers omit ids; synthesize one so the tool reply still matches.
        if (call.id.isEmpty()) call.copy(id = "call_$n") else call
    }
}