package pw.binom.llmproxy

import io.ktor.client.HttpClient
import io.ktor.client.call.body
import io.ktor.client.request.headers
import io.ktor.utils.io.LineEnding
import io.ktor.utils.io.readAvailable
import io.ktor.utils.io.readLine
import io.ktor.http.Headers
import io.ktor.client.request.preparePost
import io.ktor.client.request.setBody
import io.ktor.client.statement.HttpResponse
import io.ktor.http.ContentType
import io.ktor.http.HttpStatusCode
import io.ktor.server.application.Application
import io.ktor.server.application.ApplicationCall
import io.ktor.server.application.call
import io.ktor.server.application.install
import io.ktor.server.request.httpMethod
import io.ktor.server.request.receiveText
import io.ktor.server.request.uri
import io.ktor.server.response.respondBytesWriter
import io.ktor.server.response.respondText
import io.ktor.server.routing.get
import io.ktor.server.routing.post
import io.ktor.server.routing.routing
import io.ktor.utils.io.ByteReadChannel
import io.ktor.utils.io.ByteWriteChannel
import io.ktor.utils.io.writeFully
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.sync.Mutex
import kotlinx.datetime.Clock
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonArray
import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.JsonNull
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.jsonArray
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.jsonPrimitive
import net.mamoe.yamlkt.Yaml
import net.mamoe.yamlkt.YamlElement
import net.mamoe.yamlkt.YamlList
import net.mamoe.yamlkt.YamlLiteral
import net.mamoe.yamlkt.YamlMap
import kotlin.concurrent.Volatile
import kotlin.time.Duration.Companion.seconds
import kotlin.time.TimeSource

/**
 * Прослойка OpenAI API (YAML-декларативный роутер, Kotlin Multiplatform).
 *
 * Конфигурация целиком в YAML-файле (см. CONFIG.md): объявляются провайдеры
 * (`providers`), каталог апстрим-моделей (`upstreams`) и модели, видимые клиенту
 * (`models`). Маршрутизация между апстримами модели — по внутреннему учёту
 * конкурентности (AtomicInt на апстрим); в тело запроса послойно
 * вмердживаются `patch` (provider → upstream → model).
 *
 * Env (только эти):
 *  - CONFIG_PATH  путь к YAML; default `config.yaml` в каталоге проекта (CWD)
 *
 * Порт/интерфейс биндинга задаются блоком `server` в YAML (см. CONFIG.md).
 */
fun main() {
    val path = getEnv("CONFIG_PATH") ?: "config.yaml"
    val root = Yaml.decodeYamlFromString(readConfigText(path))
    val config = parseConfig(root)
    val providersById = config.providers.associateBy { it.id }
    val upstreamsById = config.upstreams.associateBy { it.id }

    config.upstreams.forEach { up ->
        if (up.provider !in providersById) {
            log.warn { "[llm-proxy] upstream '${up.id}' ссылается на несуществующего provider '${up.provider}'" }
        }
    }
    config.models.forEach { m ->
        m.upstreams.forEach { ref ->
            if (ref !in upstreamsById) {
                log.warn { "[llm-proxy] model '${m.name}' ссылается на несуществующий upstream '$ref'" }
            }
        }
    }

    val active = config.upstreams.associate { up ->
        up.id to UpstreamCounter(effectiveConcurrencyLimit(up, providersById[up.provider]))
    }

    log.info {
        "[llm-proxy] загружено: providers=${config.providers.size}, " +
            "upstreams=${config.upstreams.size}, models=${config.models.size} (config=$path)"
    }
    log.info {
        "[llm-proxy] server: bind host=${config.server.host} port=${config.server.port}"
    }
    log.info {
        "[llm-proxy] upstream: request_timeout=${UPSTREAM_REQUEST_TIMEOUT_MS}ms"
    }

    val http = createHttpClient()
    val sessions = SessionRegistry()
    startServer(config.server.host, config.server.port) {
        proxyModule(config, providersById, upstreamsById, active, sessions, http)
    }
}

private val json = Json { ignoreUnknownKeys = true }

fun Application.proxyModule(
    config: Config,
    providersById: Map<String, ProviderConf>,
    upstreamsById: Map<String, UpstreamConf>,
    active: Map<String, UpstreamCounter>,
    sessions: SessionRegistry,
    http: HttpClient,
) {
    routing {
        post("/v1/chat/completions") {
            handleChat(call, config, providersById, upstreamsById, active, sessions, http)
        }
        get("/v1/models") {
            handleModels(call, config)
        }
    }
}

private suspend fun handleChat(
    call: ApplicationCall,
    config: Config,
    providersById: Map<String, ProviderConf>,
    upstreamsById: Map<String, UpstreamConf>,
    active: Map<String, UpstreamCounter>,
    sessions: SessionRegistry,
    http: HttpClient,
) {
    log.info { "[llm-proxy] chat ${call.request.httpMethod.value} ${call.request.uri} headers: ${formatHeadersForLog(call.request.headers)}" }
    val raw = call.receiveText()
    if (raw.isBlank()) {
        call.respondText(errorJson("empty body"), ContentType.Application.Json, HttpStatusCode.BadRequest)
        return
    }
    val bodyJson = try {
        json.parseToJsonElement(raw).jsonObject
    } catch (e: Exception) {
        call.respondText(errorJson("bad json: ${e.message}"), ContentType.Application.Json, HttpStatusCode.BadRequest)
        return
    }

    val modelName = bodyJson["model"]?.jsonPrimitive?.content ?: "?"
    val modelConf = config.models.firstOrNull { it.name == modelName }
    if (modelConf == null) {
        call.respondText(errorJson("unknown model: $modelName"), ContentType.Application.Json, HttpStatusCode.NotFound)
        return
    }

    val pool = modelConf.upstreams.mapNotNull { upstreamsById[it] }
    if (pool.isEmpty()) {
        call.respondText(errorJson("model $modelName has no valid upstreams"), ContentType.Application.Json, HttpStatusCode.BadGateway)
        return
    }

    val clientWantsStream = bodyJson["stream"]?.jsonPrimitive?.content == "true"
    val failed = mutableSetOf<String>()
    var anyClaimed = false

    while (true) {
        val up = pickFreeUpstream(pool, active, failed) ?: break
        anyClaimed = true
        val start = TimeSource.Monotonic.markNow()
        try {
            val provider = providersById[up.provider]
            if (provider == null) {
                log.error { "[llm-proxy] upstream '${up.id}': provider '${up.provider}' не найден (конфиг)" }
                failed.add(up.id)
                continue
            }

            val patched = buildBody(bodyJson, provider, up, modelConf)
            val forwarded = if (clientWantsStream) {
                patched
            } else {
                val m = patched.toMutableMap().apply {
                    this["stream"] = JsonPrimitive(true)
                    val so = (this["stream_options"] as? JsonObject)?.toMutableMap() ?: mutableMapOf()
                    so["include_usage"] = JsonPrimitive(true)
                    this["stream_options"] = JsonObject(so)
                }
                JsonObject(m)
            }

            val url = provider.url.trimEnd('/') + "/chat/completions"
            val providerKey = resolveEnv(provider.key)
            val sessionHeader = provider.session_header
            // Клиентский одноимённый заголовок не пробрасываем: при метке
            // значение x-opencode-session всегда вычисляем сами (LCP по истории).
            val forwardedHeaders = headersToForward(call.request.headers)
                .filterKeys { sessionHeader == null || !it.equals(sessionHeader, ignoreCase = true) }
            val sessionId = sessionHeader?.let { sessions.resolve(sessionPrefixHashes(forwarded)) }
            val outgoingHeaders = forwardedHeaders.toMutableMap().apply {
                this["Content-Type"] = listOf("application/json")
                if (providerKey.isNotEmpty()) this["Authorization"] = listOf("Bearer $providerKey")
                if (sessionHeader != null && sessionId != null) this[sessionHeader] = listOf(sessionId)
            }
            log.info {
                "[llm-proxy] chat model=$modelName upstream=${up.id} session=${sessionId ?: "-"} → $url headers: ${formatHeadersForLog(outgoingHeaders)}"
            }
            var failover = false
            var responded = false
            var upstreamStatus = 0

            http.preparePost(url) {
                headers {
                    forwardedHeaders.forEach { (name, values) ->
                        appendAll(name, values)
                    }
                    // Авторизация — всегда наша (ключ провайдера из конфига);
                    // клиентский Authorization не пересылается. Если у провайдера
                    // ключ не задан — Authorization не отправляем вовсе.
                    if (providerKey.isNotEmpty()) {
                        set("Authorization", "Bearer $providerKey")
                    }
                    set("Content-Type", "application/json")
                    if (sessionHeader != null && sessionId != null) {
                        set(sessionHeader, sessionId)
                    }
                }
                setBody(forwarded.toString())
            }.execute { resp ->
                upstreamStatus = resp.status.value
                if (upstreamStatus >= 500 || upstreamStatus == 429) {
                    log.warn { "[llm-proxy] model=$modelName upstream=${up.id} вернул $upstreamStatus → фейловер" }
                    failed.add(up.id)
                    failover = true
                    return@execute
                }
                responded = true
                if (clientWantsStream) {
                    val ct = resp.headers["Content-Type"] ?: "text/event-stream"
                    val status = HttpStatusCode.fromValue(upstreamStatus)
                    val thinkMode = effectiveThinkTags(up, provider)
                    call.respondBytesWriter(ContentType.parse(ct), status) {
                        val ch = resp.body<ByteReadChannel>()
                        if (thinkMode == "off") {
                            // Флажок не выставлен — сырой байтовый passthrough как раньше.
                            val buf = ByteArray(8192)
                            while (true) {
                                val n = ch.readAvailable(buf)
                                if (n == -1) break
                                if (n > 0) {
                                    writeFully(buf, 0, n)
                                    flush()
                                }
                            }
                        } else {
                            streamSseWithThinkTags(ch, thinkMode)
                        }
                    }
                    log.info {
                        "[llm-proxy] chat model=$modelName upstream=${up.id} provider=${provider.id} " +
                            "status=$upstreamStatus за ${start.elapsedNow().inWholeMilliseconds}ms stream=true"
                    }
                } else {
                    val ct = resp.headers["Content-Type"] ?: "application/json"
                    val full = resp.body<String>()
                    val thinkMode = effectiveThinkTags(up, provider)
                    val rebuilt = if (ct.contains("text/event-stream")) rebuildFromChunks(full) else full
                    val out = if (thinkMode != "off") (
                        runCatching { transformThinkMessage(json.parseToJsonElement(rebuilt).jsonObject, thinkMode) }
                            .getOrNull() ?: rebuilt
                    ) else rebuilt
                    val outCt = runCatching {
                        if (ct.contains("text/event-stream") &&
                            Json.parseToJsonElement(out).jsonObject["error"] != null
                        ) {
                            ContentType.Application.Json
                        } else {
                            ContentType.parse(ct)
                        }
                    }.getOrDefault(ContentType.parse(ct))
                    call.respondText(out, outCt, HttpStatusCode.fromValue(upstreamStatus))
                    log.info {
                        "[llm-proxy] chat model=$modelName upstream=${up.id} provider=${provider.id} " +
                            "status=$upstreamStatus за ${start.elapsedNow().inWholeMilliseconds}ms stream=false"
                    }
                }
            }

            if (failover) continue
            if (responded) return
        } catch (e: CancellationException) {
            log.info { "[llm-proxy] chat model=$modelName upstream=${up.id} ОТМЕНЕНО клиентом за ${start.elapsedNow().inWholeMilliseconds}ms" }
            throw e
        } catch (e: Exception) {
            log.error { "[llm-proxy] chat model=$modelName upstream=${up.id} ОШИБКА: ${e.message} за ${start.elapsedNow().inWholeMilliseconds}ms → фейловер" }
            failed.add(up.id)
            continue
        } finally {
            release(up, active)
        }
    }

    if (anyClaimed) {
        call.respondText(errorJson("all upstreams failed"), ContentType.Application.Json, HttpStatusCode.ServiceUnavailable)
    } else {
        call.respondText(errorJson("all upstreams busy"), ContentType.Application.Json, HttpStatusCode.ServiceUnavailable)
    }
}

private fun errorJson(msg: String): String =
    """{"error":{"message":"$msg"}}"""

/** Hop-by-hop и прокси-специфичные заголовки, которые НЕ пересылаются в апстрим. */
private val SKIP_HEADER_NAMES = setOf(
    "host",
    "content-length",
    "transfer-encoding",
    "te",
    "connection",
    "proxy-connection",
    "keep-alive",
    "upgrade",
    // Authorization управляется прокси явно (ключ провайдера), клиентский
    // не пересылается
    "authorization",
    // Content-Type всегда наш (application/json: тело мержится как JSON),
    // клиентский не пересылаем, чтобы не ушло двух заголовков
    "content-type",
)

/**
 * Заголовки входящего запроса для пересылки в апстрим: все, кроме служебных
 * ([SKIP_HEADER_NAMES]). Тело может быть изменено `patch`, а соединение до
 * апстрима — другое, поэтому Content-Length/Transfer-Encoding/Host/Connection
 * управляет сам прокси (клиентский Ktor выставит свои автоматически).
 * `Authorization` тоже в списке пропусков — прокси управляет им явно
 * (ключ провайдера из конфига; клиентский не пересылается).
 */
internal fun headersToForward(request: Headers): Map<String, List<String>> =
    request.entries()
        .filter { (name, _) -> name.lowercase() !in SKIP_HEADER_NAMES }
        .associate { (name, values) -> name to values }

/** Заголовки, значения которых маскируются в логах (секреты клиента). */
private val SENSITIVE_HEADER_NAMES = setOf(
    "authorization",
    "proxy-authorization",
    "x-api-key",
    "api-key",
    "cookie",
    "set-cookie",
)

/**
 * Заголовки в виде строки для лога (`name=v1|v2, ...`). Значения чувствительных
 * имён ([SENSITIVE_HEADER_NAMES]) маскируются `***`, чтобы не светить секреты.
 */
internal fun formatHeadersForLog(headers: Map<String, List<String>>): String =
    headers.entries.joinToString(", ") { (name, values) ->
        val shown = if (name.lowercase() in SENSITIVE_HEADER_NAMES) values.map { "***" } else values
        "$name=${shown.joinToString("|")}"
    }

internal fun formatHeadersForLog(headers: Headers): String =
    formatHeadersForLog(headers.entries().associate { (name, values) -> name to values })

/**
 * Сборка тела запроса: подмена `model` на реальное имя апстрима + глубокий
 * послойный мерж `patch` в порядке provider → upstream → model.
 */
internal fun buildBody(base: JsonObject, provider: ProviderConf, up: UpstreamConf, model: ModelConf): JsonObject {
    val withModel = base.toMutableMap().apply { this["model"] = JsonPrimitive(up.model) }
    var acc = JsonObject(withModel)
    listOf(provider.patch, up.patch, model.patch).filterNotNull().forEach { patch ->
        acc = merge(acc, patch)
    }
    return acc
}

/** Конвертация YAML-дерева (YamlElement) в kotlinx JsonElement для мержа в тело. */
internal fun yamlToJson(el: YamlElement): JsonElement = when (el) {
    is YamlMap -> JsonObject(el.mapKeys { (k, _) -> yamlKey(k) }.mapValues { (_, v) -> yamlToJson(v) })
    is YamlList -> JsonArray(el.map { yamlToJson(it) })
    is YamlLiteral -> yamlLiteralToJson(el)
    else -> JsonNull
}

internal fun yamlKey(key: YamlElement): String =
    if (key is YamlLiteral) key.content else key.toString()

internal fun yamlLiteralToJson(lit: YamlLiteral): JsonPrimitive {
    val c = lit.content
    val bool = c.toBooleanStrictOrNull()
    val lng = c.toLongOrNull()
    val dbl = c.toDoubleOrNull()
    return when {
        bool != null -> JsonPrimitive(bool)
        lng != null -> JsonPrimitive(lng)
        dbl != null -> JsonPrimitive(dbl)
        else -> JsonPrimitive(c)
    }
}

/** Глубокий мерж: вложенные объекты сливаются рекурсивно, скаляры/массивы заменяются. */
internal fun merge(base: JsonObject, patch: JsonObject): JsonObject {
    val merged = base.toMutableMap()
    for ((k, v) in patch) {
        merged[k] = when {
            v is JsonObject && merged[k] is JsonObject ->
                merge(merged[k] as JsonObject, v)
            else -> v
        }
    }
    return JsonObject(merged)
}

/** Подстановка `${ENV}` в строках ключей/url из переменных среды. */
internal fun resolveEnv(s: String): String =
    """\$\{([^}]+)\}""".toRegex().replace(s) { m -> getEnv(m.groupValues[1]) ?: "" }

/** Атомарно занять слот у апстрима (по эффективному лимиту); false, если все заняты. */
internal fun tryClaim(up: UpstreamConf, active: Map<String, UpstreamCounter>): Boolean =
    active.getValue(up.id).tryClaim()

/**
 * Эффективный лимит конкурентности апстрима: значение у апстрима (модели), если
 * задано; иначе у провайдера; иначе безлимит.
 */
internal fun effectiveConcurrencyLimit(up: UpstreamConf, provider: ProviderConf?): Int =
    up.max_concurrency ?: provider?.max_concurrency ?: Int.MAX_VALUE

/**
 * Эффективные think_tags апстрима: значение у апстрима, если задано; иначе у
 * провайдера; иначе "off". Значения "true" трактуются как "split", "false" и
 * любое неизвестное/пустое — как "off".
 */
internal fun effectiveThinkTags(up: UpstreamConf, provider: ProviderConf?): String {
    val raw = up.think_tags ?: provider?.think_tags ?: "off"
    return when (raw) {
        "split", "strip" -> raw
        "true" -> "split"
        else -> "off"
    }
}

/** Освободить слот апстрима (в finally по завершении проксирования). */
internal fun release(up: UpstreamConf, active: Map<String, UpstreamCounter>) {
    active.getValue(up.id).release()
}

class UpstreamCounter(private val limit: Int, initial: Int = 0) {
    private val lock = Mutex()
    @Volatile
    private var count: Int = initial

    fun tryClaim(): Boolean {
        if (!lock.tryLock()) return false
        try {
            return if (count >= limit) false else {
                count++
                true
            }
        } finally {
            lock.unlock()
        }
    }

    fun release(): Unit {
        if (!lock.tryLock()) return
        try {
            if (count > 0) count--
        } finally {
            lock.unlock()
        }
    }

    val current: Int
        get() {
            if (!lock.tryLock()) return -1
            try {
                return count
            } finally {
                lock.unlock()
            }
        }
}

/**
 * Выбор апстрима для попытки: первый по порядку (приоритету) апстрим из `pool`,
 * у которого свободен слот и который ещё не в `excluded` (не упал ранее).
 * Сразу занимает слот (через [tryClaim]). Если свободных нет — возвращает null.
 */
internal fun pickFreeUpstream(
    pool: List<UpstreamConf>,
    active: Map<String, UpstreamCounter>,
    excluded: Set<String>,
): UpstreamConf? =
    pool.firstOrNull { up -> up.id !in excluded && tryClaim(up, active) }

private suspend fun handleModels(call: ApplicationCall, config: Config) {
    log.info { "[llm-proxy] models ${call.request.httpMethod.value} ${call.request.uri} headers: ${formatHeadersForLog(call.request.headers)}" }
    val created = TimeSource.Monotonic.markNow().elapsedNow().inWholeSeconds
    val data = config.models.map { m ->
        JsonObject(
            mapOf(
                "id" to JsonPrimitive(m.name),
                "object" to JsonPrimitive("model"),
                "created" to JsonPrimitive(created),
                "owned_by" to JsonPrimitive("llm-proxy"),
            ),
        )
    }
    val body = JsonObject(mapOf("object" to JsonPrimitive("list"), "data" to JsonArray(data)))
    call.respondText(body.toString(), ContentType.Application.Json, HttpStatusCode.OK)
}

/**
 * Собрать полный chat.completion из SSE-чанков апстрима (для non-stream клиентов).
 * Сохраняются все значимые поля, которые отдаёт апстрим: id/created/model/
 * system_fingerprint/service_tier, usage, несколько choices (по index) с
 * role/content/reasoning_content/tool_calls/finish_reason/logprobs.
 */
internal fun rebuildFromChunks(sse: String): String {
    var id = ""
    var created: Long? = null
    var model = ""
    var systemFingerprint: String? = null
    var serviceTier: String? = null
    var usage: JsonObject? = null
    var provider: JsonElement? = null
    var error: JsonObject? = null

    class MutableChoice {
        var role: String? = null
        val content = StringBuilder()
        val reasoning = StringBuilder()
        var finishReason: String? = null
        var logprobs: JsonElement? = null
        val toolCalls = mutableListOf<JsonObject>()
    }

    val choices = mutableMapOf<Int, MutableChoice>()

    sse.lineSequence().forEach { line ->
        if (!line.startsWith("data:")) return@forEach
        val data = line.removePrefix("data:").trim()
        if (data.isEmpty() || data == "[DONE]") return@forEach
        val obj = runCatching { json.parseToJsonElement(data).jsonObject }.getOrNull() ?: return@forEach
        if (id.isEmpty()) id = obj["id"]?.jsonPrimitive?.content ?: ""
        if (created == null) created = obj["created"]?.jsonPrimitive?.content?.toLongOrNull()
        if (model.isEmpty()) model = obj["model"]?.jsonPrimitive?.content ?: ""
        if (systemFingerprint == null) systemFingerprint = obj["system_fingerprint"]?.jsonPrimitive?.content
        if (serviceTier == null) serviceTier = obj["service_tier"]?.jsonPrimitive?.content
        if (provider == null) provider = obj["provider"]
        (obj["error"] as? JsonObject)?.let { error = it }
        (obj["usage"] as? JsonObject)?.let { usage = it }
        val chArr = obj["choices"]?.jsonArray ?: return@forEach
        for (ch in chArr) {
            val c = ch.jsonObject
            val idx = c["index"]?.jsonPrimitive?.content?.toIntOrNull() ?: 0
            val mc = choices.getOrPut(idx) { MutableChoice() }
            val delta = c["delta"]?.jsonObject
            if (delta != null) {
                if (mc.role == null) mc.role = delta["role"]?.jsonPrimitive?.content
                delta["content"]?.jsonPrimitive?.content?.takeIf { it != "null" }?.let { mc.content.append(it) }
                delta["reasoning_content"]?.jsonPrimitive?.content?.takeIf { it != "null" }?.let { mc.reasoning.append(it) }
                delta["tool_calls"]?.jsonArray?.forEach { tc -> (tc as? JsonObject)?.let { mc.toolCalls.add(it) } }
            }
            c["finish_reason"]?.jsonPrimitive?.content?.takeIf { it.isNotEmpty() && it != "null" }?.let { mc.finishReason = it }
            c["logprobs"]?.let { mc.logprobs = it }
        }
    }

    // Апстрим вернул ошибку (в SSE она приходит чанком data: {"error":{...}}).
    // Не «проглатываем» — отдаём клиенту как есть.
    if (error != null) {
        return JsonObject(mapOf("error" to error)).toString()
    }

    val choiceArr = choices.entries.sortedBy { it.key }.map { (idx, mc) ->
        val msg = mutableMapOf<String, JsonElement>(
            "role" to JsonPrimitive(mc.role ?: "assistant"),
            "content" to JsonPrimitive(mc.content.toString()),
        )
        if (mc.reasoning.isNotEmpty()) msg["reasoning_content"] = JsonPrimitive(mc.reasoning.toString())
        if (mc.toolCalls.isNotEmpty()) msg["tool_calls"] = JsonArray(mc.toolCalls)
        val ch = mutableMapOf<String, JsonElement>(
            "index" to JsonPrimitive(idx),
            "message" to JsonObject(msg),
            "finish_reason" to JsonPrimitive(mc.finishReason ?: "stop"),
        )
        mc.logprobs?.let { ch["logprobs"] = it }
        JsonObject(ch)
    }

    val root = mutableMapOf<String, JsonElement>(
        "id" to JsonPrimitive(id),
        "object" to JsonPrimitive("chat.completion"),
        "created" to JsonPrimitive(created ?: Clock.System.now().epochSeconds),
        "model" to JsonPrimitive(model),
        "choices" to JsonArray(choiceArr),
    )
    systemFingerprint?.let { root["system_fingerprint"] = JsonPrimitive(it) }
    serviceTier?.let { root["service_tier"] = JsonPrimitive(it) }
    provider?.let { root["provider"] = it }
    usage?.let { root["usage"] = it }
    return JsonObject(root).toString()
}

/**
 * Пересобрать полный chat.completion (обычный JSON или результат
 * [rebuildFromChunks]): по каждому choice взять `message.content` — только если
 * это JSON-строка (массив частей не трогаем) — и прогнать целиком через
 * [ThinkTagSplitter] (feed + finish). Остаток возвращается в `message.content`,
 * вырезанное — в `message.reasoning_content` (режим split; при strip не
 * добавляем). Если `reasoning_content` уже был непустой строкой — новое
 * дописываем в конец существующего, не теряя прежнее. Ни один choice не
 * изменился → null (отдать исходную строку как есть).
 */
internal fun transformThinkMessage(obj: JsonObject, thinkMode: String): String? {
    val choices = obj["choices"]?.jsonArray ?: return null
    val addReasoning = thinkMode == "split"
    var changed = false
    val newChoices = choices.map { choiceEl ->
        val choice = choiceEl.jsonObject
        val message = choice["message"]?.jsonObject ?: return@map choiceEl
        val contentStr = (message["content"] as? JsonPrimitive)?.takeIf { it.isString }?.content
        if (contentStr == null) return@map choiceEl
        val splitter = ThinkTagSplitter(thinkMode)
        val (outContent, outReasoning) = splitter.feed(contentStr)
        val (tailContent, tailReasoning) = splitter.finish()
        val reasoning = outReasoning + tailReasoning
        changed = true
        val newMessage = message.toMutableMap().apply {
            this["content"] = JsonPrimitive(outContent + tailContent)
            if (addReasoning && reasoning.isNotEmpty()) {
                val existing = (this["reasoning_content"] as? JsonPrimitive)?.takeIf { it.isString }?.content
                this["reasoning_content"] = JsonPrimitive((existing ?: "") + reasoning)
            }
        }
        JsonObject(choice.toMutableMap().apply { this["message"] = JsonObject(newMessage) })
    }
    if (!changed) return null
    return JsonObject(obj.toMutableMap().apply { this["choices"] = JsonArray(newChoices) }).toString()
}

/**
 * Построчный разбор SSE-стрима с рассечением think-тегов. Строки, не начинающиеся
 * с `data:`, и `data: [DONE]` уходят клиенту без изменений (с `\n`). Прочие
 * `data:`-строки парсятся и прогоняются через [transformThinkChunk]; результат
 * записывается как `data: <json>\n\n` (событие-граница SSE), а при ошибке парса —
 * исходная строка. Каждую строку сразу `flush()`, чтобы стрим не «залипал» в
 * буфере. В конце потока накопленные хвосты сплиттеров сбрасываются финиш-чанком.
 */
internal suspend fun ByteWriteChannel.streamSseWithThinkTags(source: ByteReadChannel, thinkMode: String) {
    val splitters = mutableMapOf<Int, ThinkTagSplitter>()
    val addReasoning = thinkMode == "split"
    while (true) {
        val line = source.readLine(LineEnding.Lenient) ?: break
        when {
            line.startsWith("data:") -> {
                val payload = line.removePrefix("data:").trim()
                if (payload == "[DONE]") {
                    emitUtf8("data: [DONE]\n\n")
                } else {
                    val obj = runCatching { json.parseToJsonElement(payload).jsonObject }.getOrNull()
                    val out = obj?.let { transformThinkChunk(it, splitters, thinkMode, addReasoning) }
                    if (out == null) emitUtf8("$line\n\n") else emitUtf8("data: $out\n\n")
                }
            }
            else -> emitUtf8("$line\n")
        }
        flush()
    }
    if (thinkMode == "split" || thinkMode == "strip") {
        splitters.forEach { (idx, sp) ->
            val (tail, reasoning) = sp.finish()
            val hasContent = tail.isNotEmpty()
            val hasReasoning = addReasoning && reasoning.isNotEmpty()
            if (!hasContent && !hasReasoning) return@forEach
            val delta = mutableMapOf<String, JsonElement>()
            if (hasContent) delta["content"] = JsonPrimitive(tail)
            if (hasReasoning) delta["reasoning_content"] = JsonPrimitive(reasoning)
            val chunk = JsonObject(
                mutableMapOf(
                    "choices" to JsonArray(
                        listOf(
                            JsonObject(
                                mutableMapOf(
                                    "index" to JsonPrimitive(idx),
                                    "delta" to JsonObject(delta),
                                ),
                            ),
                        ),
                    ),
                ),
            )
            emitUtf8("data: $chunk\n\n")
            flush()
        }
    }
}

/**
 * Пересобрать SSE-чанк: по каждому choice (ключ `index`, дефолт 0) взять
 * `delta.content` (только если это JSON-строка; массив частей не трогаем) и
 * прогнать через [ThinkTagSplitter] для этого index. Остаток возвращается в
 * `delta.content` (поле убирается, если пустое); вырезанное — в
 * `delta.reasoning_content` (только режим split, при strip не добавляем).
 * Чанк без `choices` или без строкового `delta.content` не меняется —
 * возвращается null (отдать исходную строку как есть).
 */
internal fun transformThinkChunk(
    obj: JsonObject,
    splitters: MutableMap<Int, ThinkTagSplitter>,
    thinkMode: String,
    addReasoning: Boolean,
): String? {
    val choices = obj["choices"]?.jsonArray ?: return null
    var changed = false
    val newChoices = choices.map { choiceEl ->
        val choice = choiceEl.jsonObject
        val delta = choice["delta"]?.jsonObject
        val contentStr = (delta?.get("content") as? JsonPrimitive)?.takeIf { it.isString }?.content
        if (contentStr == null) return@map choiceEl
        val idx = choice["index"]?.jsonPrimitive?.content?.toIntOrNull() ?: 0
        val splitter = splitters.getOrPut(idx) { ThinkTagSplitter(thinkMode) }
        val (newContent, reasoning) = splitter.feed(contentStr)
        changed = true
        val newDelta = delta.toMutableMap()
        if (newContent.isEmpty()) newDelta.remove("content") else newDelta["content"] = JsonPrimitive(newContent)
        if (addReasoning && reasoning.isNotEmpty()) newDelta["reasoning_content"] = JsonPrimitive(reasoning)
        JsonObject(choice.toMutableMap().apply { this["delta"] = JsonObject(newDelta) })
    }
    if (!changed) return null
    return JsonObject(obj.toMutableMap().apply { this["choices"] = JsonArray(newChoices) }).toString()
}

/** Записать строку как UTF-8 байты (KMP-безопасно, без java.io). */
private suspend fun ByteWriteChannel.emitUtf8(text: String) {
    val bytes = text.encodeToByteArray()
    writeFully(bytes, 0, bytes.size)
}

@Serializable
data class ProviderConf(
    val id: String,
    val url: String,
    val key: String = "",
    val max_concurrency: Int? = null,
    val patch: JsonObject? = null,
    val session_header: String? = null,
    val think_tags: String? = null,
)

data class UpstreamConf(
    val id: String,
    val provider: String,
    val model: String,
    val max_concurrency: Int? = null,
    val patch: JsonObject? = null,
    val think_tags: String? = null,
)

data class ModelConf(
    val name: String,
    val upstreams: List<String>,
    val patch: JsonObject? = null,
)

data class ServerConf(
    val host: String = "0.0.0.0",
    val port: Int = 8100,
)

data class Config(
    val server: ServerConf,
    val providers: List<ProviderConf>,
    val upstreams: List<UpstreamConf>,
    val models: List<ModelConf>,
)

/**
 * Ручной разбор YAML-дерева (YamlElement) в конфиг. Поле `patch` сразу
 * конвертируется в kotlinx JsonObject (merge`ить в тело тривиально).
 * Декодирование всего документа как YamlElement надёжнее, чем вложенный
 * `@Serializable`-класс с полем YamlElement (yamlkt криво читает списки).
 */
internal fun parseConfig(root: YamlElement): Config {
    val rootMap = root as? YamlMap ?: error("config root must be a map")
    val top = rootMap.toMap()

    fun list(key: String): List<YamlElement> {
        val v = top[key] ?: return emptyList()
        return (v as? YamlList)?.map { it } ?: emptyList()
    }

    val serverMap = (top["server"] as? YamlMap)?.toMap()
    val server = if (serverMap != null) {
        ServerConf(
            host = serverMap.strOrNull("host") ?: "0.0.0.0",
            port = serverMap.strOrNull("port")?.toIntOrNull() ?: 8100,
        )
    } else {
        ServerConf()
    }

    val providers = list("providers").map { entry ->
        val m = (entry as YamlMap).toMap()
        ProviderConf(
            id = m.str("id"),
            url = m.str("url"),
            key = m.strOrNull("key") ?: "",
            max_concurrency = m.strOrNull("max_concurrency")?.toIntOrNull(),
            patch = m.yamlMapOrNull("patch")?.let { yamlToJson(it) as JsonObject },
            session_header = m.strOrNull("session_header"),
            think_tags = m.strOrNull("think_tags"),
        )
    }

    val upstreams = list("upstreams").map { entry ->
        val m = (entry as YamlMap).toMap()
        UpstreamConf(
            id = m.str("id"),
            provider = m.str("provider"),
            model = m.str("model"),
            max_concurrency = m.strOrNull("max_concurrency")?.toIntOrNull(),
            patch = m.yamlMapOrNull("patch")?.let { yamlToJson(it) as JsonObject },
            think_tags = m.strOrNull("think_tags"),
        )
    }

    val models = list("models").map { entry ->
        val m = (entry as YamlMap).toMap()
        ModelConf(
            name = m.str("name"),
            upstreams = (m["upstreams"] as? YamlList)?.map { (it as YamlLiteral).content } ?: emptyList(),
            patch = m.yamlMapOrNull("patch")?.let { yamlToJson(it) as JsonObject },
        )
    }

    return Config(server, providers, upstreams, models)
}

/** YamlMap -> Map<String, YamlElement> (ключи YAML — строковые скаляры). */
internal fun YamlMap.toMap(): Map<String, YamlElement> =
    entries.associate { (it.key as YamlLiteral).content to it.value }

internal fun Map<String, YamlElement>.str(key: String): String =
    (this[key] as? YamlLiteral)?.content ?: error("config: missing string '$key'")

internal fun Map<String, YamlElement>.strOrNull(key: String): String? =
    (this[key] as? YamlLiteral)?.content

internal fun Map<String, YamlElement>.yamlMapOrNull(key: String): YamlMap? =
    this[key] as? YamlMap