llm-proxy/src/commonTest/kotlin/pw/binom/llmproxy/ConfigLogicTest.kt

package pw.binom.llmproxy

import kotlin.test.Test
import kotlin.test.assertEquals
import kotlin.test.assertFalse
import kotlin.test.assertTrue
import io.ktor.http.headersOf
import kotlinx.coroutines.test.runTest
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.jsonArray
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.jsonPrimitive
import net.mamoe.yamlkt.Yaml

class ConfigLogicTest {

    private val noBackoff = BackoffRegistry(emptyMap(), emptyMap())

    @Test
    fun mergeDeepMergesNestedObjectsAndReplacesScalars() {
        val base = Json.parseToJsonElement("""{"a":{"x":1,"y":2},"b":1}""").jsonObject
        val patch = Json.parseToJsonElement("""{"a":{"y":3,"z":4},"c":5}""").jsonObject
        val merged = merge(base, patch)
        assertEquals("""{"a":{"x":1,"y":3,"z":4},"b":1,"c":5}""", merged.toString())
    }

    @Test
    fun mergeReplacesArraysByKey() {
        val base = Json.parseToJsonElement("""{"provider":{"ignore":["a"]}}""").jsonObject
        val patch = Json.parseToJsonElement("""{"provider":{"ignore":["b","c"]}}""").jsonObject
        assertEquals("""{"provider":{"ignore":["b","c"]}}""", merge(base, patch).toString())
    }

    @Test
    fun parseConfigReadsAllBlocksAndConvertsPatchToJsonObject() {
        val yaml = """
            providers:
              - id: p1
                url: "https://x.ru/api/v1"
                key: "k"
                max_concurrency: 4
                patch:
                  provider:
                    allow_fallbacks: false
            upstreams:
              - id: u1
                provider: p1
                model: real-1
                max_concurrency: 2
                patch:
                  provider:
                    ignore: [deepseek]
              - id: bad
                provider: missing
                model: x
            models:
              - name: m1
                upstreams: [u1, bad]
                patch:
                  reasoning:
                    enabled: false
        """.trimIndent()

        val cfg = parseConfig(Yaml.decodeYamlFromString(yaml))

        assertEquals(1, cfg.providers.size)
        assertEquals("https://x.ru/api/v1", cfg.providers[0].url)
        assertEquals(4, cfg.providers[0].max_concurrency)
        assertEquals(2, cfg.upstreams.size)
        assertEquals(2, cfg.upstreams[0].max_concurrency)
        assertEquals(listOf("u1", "bad"), cfg.models[0].upstreams)

        assertEquals("""{"provider":{"allow_fallbacks":false}}""", cfg.providers[0].patch.toString())
        assertEquals("""{"provider":{"ignore":["deepseek"]}}""", cfg.upstreams[0].patch.toString())
        assertEquals("""{"reasoning":{"enabled":false}}""", cfg.models[0].patch.toString())
    }

    @Test
    fun parseConfigReadsServerBlock() {
        val yaml = """
            server:
              host: 127.0.0.1
              port: 9200
            models:
              - name: m1
                upstreams: []
        """.trimIndent()

        val cfg = parseConfig(Yaml.decodeYamlFromString(yaml))

        assertEquals("127.0.0.1", cfg.server.host)
        assertEquals(9200, cfg.server.port)
    }

    @Test
    fun parseConfigServerDefaultsWhenBlockAbsent() {
        val yaml = """
            models:
              - name: m1
                upstreams: []
        """.trimIndent()

        val cfg = parseConfig(Yaml.decodeYamlFromString(yaml))

        assertEquals("0.0.0.0", cfg.server.host)
        assertEquals(8100, cfg.server.port)
    }

    @Test
    fun parseConfigServerFieldDefaultsPerField() {
        val yaml = """
            server:
              host: 192.168.88.10
            models:
              - name: m1
                upstreams: []
        """.trimIndent()

        val cfg = parseConfig(Yaml.decodeYamlFromString(yaml))

        assertEquals("192.168.88.10", cfg.server.host)
        assertEquals(8100, cfg.server.port)
    }

    @Test
    fun parseConfigOmitsPatchWhenAbsent() {
        val yaml = """
            providers:
              - id: p1
                url: "https://x.ru/api/v1"
            upstreams:
              - id: u1
                provider: p1
                model: real-1
            models:
              - name: m1
                upstreams: [u1]
        """.trimIndent()
        val cfg = parseConfig(Yaml.decodeYamlFromString(yaml))
        assertEquals(null, cfg.providers[0].patch)
        assertEquals(null, cfg.models[0].patch)
        assertEquals(null, cfg.upstreams[0].max_concurrency)
    }

    @Test
    fun buildBodySubstitutesModelAndAppliesLayersInOrder() {
        val provider = ProviderConf(
            "p1", "https://x", "", null,
            Json.parseToJsonElement("""{"provider":{"allow_fallbacks":false}}""").jsonObject,
        )
        val up = UpstreamConf(
            "u1", "p1", "real-1", null,
            Json.parseToJsonElement("""{"provider":{"ignore":["deepseek"]}}""").jsonObject,
        )
        val model = ModelConf(
            "m1", listOf("u1"),
            Json.parseToJsonElement("""{"reasoning":{"enabled":false}}""").jsonObject,
        )
        val body = Json.parseToJsonElement(
            """{"model":"m1","messages":[],"temperature":0.7}""",
        ).jsonObject

        val out = buildBody(body, provider, up, model)
        assertEquals(
            """{"model":"real-1","messages":[],"temperature":0.7,"provider":{"allow_fallbacks":false,"ignore":["deepseek"]},"reasoning":{"enabled":false}}""",
            out.toString(),
        )
    }

    @Test
    fun buildBodyWithNoPatchesReturnsClientBodyWithSubstitutedModel() {
        val provider = ProviderConf("p1", "https://x", "")
        val up = UpstreamConf("u1", "p1", "real-1", null, null)
        val model = ModelConf("m1", listOf("u1"), null)
        val body = Json.parseToJsonElement("""{"model":"m1","a":1}""").jsonObject
        assertEquals("""{"model":"real-1","a":1}""", buildBody(body, provider, up, model).toString())
    }

    @Test
    fun resolveEnvSubstitutesMissingVarWithEmptyAndPassesLiteralThrough() {
        assertEquals("", resolveEnv("\${LLM_PROXY_TEST_MISSING_VAR}"))
        assertEquals("plain", resolveEnv("plain"))
        assertEquals("pre--post", resolveEnv("pre-\${LLM_PROXY_TEST_MISSING_VAR}-post"))
    }

    @Test
    fun tryClaimRespectsMaxConcurrencyAndReleaseFreesSlot() {
        val active = mapOf("u1" to UpstreamCounter(1))
        val up = UpstreamConf("u1", "p", "m", 1, null)
        assertTrue(tryClaim(up, active))
        assertFalse(tryClaim(up, active))
        release(up, active)
        assertTrue(tryClaim(up, active))
    }

    @Test
    fun tryClaimUnlimitedWhenMaxConcurrencyIsNull() {
        val active = mapOf("u2" to UpstreamCounter(Int.MAX_VALUE))
        val up = UpstreamConf("u2", "p", "m", null, null)
        assertTrue(tryClaim(up, active))
        assertTrue(tryClaim(up, active))
    }

    @Test
    fun effectiveConcurrencyLimitPrefersUpstreamThenProvider() {
        val providerWithLimit = ProviderConf("p1", "https://x", "", 3, null)
        val providerWithoutLimit = ProviderConf("p2", "https://x", "", null, null)
        assertEquals(
            2,
            effectiveConcurrencyLimit(UpstreamConf("u1", "p1", "m", 2, null), providerWithLimit),
        )
        assertEquals(
            3,
            effectiveConcurrencyLimit(UpstreamConf("u1", "p1", "m", null, null), providerWithLimit),
        )
        assertEquals(
            Int.MAX_VALUE,
            effectiveConcurrencyLimit(UpstreamConf("u1", "p2", "m", null, null), providerWithoutLimit),
        )
        assertEquals(
            Int.MAX_VALUE,
            effectiveConcurrencyLimit(UpstreamConf("u1", "missing", "m", null, null), null),
        )
    }

    @Test
    fun pickFreeUpstreamReturnsFirstFreeInDeclarationOrder() = runTest {
        val active = mapOf("u1" to UpstreamCounter(1), "u2" to UpstreamCounter(2))
        val pool = listOf(
            UpstreamConf("u1", "p", "m", 1, null),
            UpstreamConf("u2", "p", "m", 2, null),
        )
        val up = pickFreeUpstream(pool, active, emptySet(), noBackoff)
        assertEquals("u1", up?.id)
        // слот реально занят
        assertEquals(1, active.getValue("u1").current)
    }

    @Test
    fun pickFreeUpstreamSkipsExcluded() = runTest {
        val active = mapOf("u1" to UpstreamCounter(1), "u2" to UpstreamCounter(2))
        val pool = listOf(
            UpstreamConf("u1", "p", "m", 1, null),
            UpstreamConf("u2", "p", "m", 2, null),
        )
        val up = pickFreeUpstream(pool, active, setOf("u1"), noBackoff)
        assertEquals("u2", up?.id)
    }

    @Test
    fun pickFreeUpstreamReturnsNullWhenAllBusy() = runTest {
        val active = mapOf("u1" to UpstreamCounter(1, 1)) // уже на лимите 1
        val pool = listOf(UpstreamConf("u1", "p", "m", 1, null))
        assertEquals(null, pickFreeUpstream(pool, active, emptySet(), noBackoff))
    }

    @Test
    fun pickFreeUpstreamImplementsFailoverOrder() = runTest {
        // dead исключён (упал ранее) — выбирается следующий живой u1
        val active = mapOf("dead" to UpstreamCounter(1), "u1" to UpstreamCounter(1))
        val pool = listOf(
            UpstreamConf("dead", "p", "m", 1, null),
            UpstreamConf("u1", "p", "m", 1, null),
        )
        val up = pickFreeUpstream(pool, active, setOf("dead"), noBackoff)
        assertEquals("u1", up?.id)
    }

    @Test
    fun pickFreeUpstreamExhaustsConcurrencyThenReturnsNull() = runTest {
        val active = mapOf("u1" to UpstreamCounter(1), "u2" to UpstreamCounter(1))
        val pool = listOf(
            UpstreamConf("u1", "p", "m", 1, null),
            UpstreamConf("u2", "p", "m", 1, null),
        )
        assertEquals("u1", pickFreeUpstream(pool, active, emptySet(), noBackoff)?.id)
        assertEquals("u2", pickFreeUpstream(pool, active, emptySet(), noBackoff)?.id)
        assertEquals(null, pickFreeUpstream(pool, active, emptySet(), noBackoff)?.id)
    }

    @Test
    fun headersToForwardKeepsUnknownAndDropsServiceHeaders() {
        val req = headersOf(
            "X-Opencode-Session" to listOf("abc-123"),
            "X-Custom" to listOf("a", "b"),
            "Host" to listOf("proxy:8100"),
            "Content-Length" to listOf("42"),
            "Transfer-Encoding" to listOf("chunked"),
            "Connection" to listOf("keep-alive"),
            "TE" to listOf("trailers"),
            "Proxy-Connection" to listOf("keep-alive"),
            "Upgrade" to listOf("h2c"),
            "Authorization" to listOf("Bearer client-secret"),
            "Content-Type" to listOf("application/x-www-form-urlencoded"),
            "Accept" to listOf("*/*"),
        )
        val out = headersToForward(req)
        assertEquals(listOf("abc-123"), out["X-Opencode-Session"])
        assertEquals(listOf("a", "b"), out["X-Custom"])
        assertEquals(listOf("*/*"), out["Accept"])
        assertEquals(3, out.size)
        assertEquals(null, out["Authorization"])
        assertEquals(null, out["Content-Type"])
    }

    @Test
    fun headersToForwardIsCaseInsensitiveOnSkipSet() {
        val out = headersToForward(
            headersOf(
                "HOST" to listOf("x"),
                "Content-length" to listOf("1"),
                "x-opencode-session" to listOf("s"),
            ),
        )
        assertEquals(1, out.size)
        assertEquals(listOf("s"), out["x-opencode-session"])
    }

    @Test
    fun formatHeadersForLogMasksSecretsAndKeepsOthers() {
        val line = formatHeadersForLog(
            headersOf(
                "X-Opencode-Session" to listOf("abc-123"),
                "Authorization" to listOf("Bearer super-secret"),
                "x-api-key" to listOf("key-1"),
                "Cookie" to listOf("session=deadbeef"),
            ),
        )
        assertTrue(line.contains("X-Opencode-Session=abc-123"))
        assertTrue(line.contains("Authorization=***"))
        assertTrue(line.contains("x-api-key=***"))
        assertTrue(line.contains("Cookie=***"))
        assertFalse(line.contains("super-secret"))
        assertFalse(line.contains("deadbeef"))
    }

    @Test
    fun formatHeadersForLogJoinsMultipleValues() {
        val line = formatHeadersForLog(
            headersOf("X-Custom" to listOf("a", "b")),
        )
        assertEquals("X-Custom=a|b", line)
    }

    @Test
    fun rebuildFromChunksPreservesAllUpstreamFields() {
        val sse = """
            data: {"id":"c1","object":"chat.completion.chunk","created":123,"model":"m","provider":"Baidu","system_fingerprint":"sf1","choices":[{"index":0,"delta":{"role":"assistant","content":"He"},"finish_reason":null}]}

            data: {"id":"c1","created":123,"model":"m","provider":"Baidu","choices":[{"index":0,"delta":{"content":"llo"},"finish_reason":null}]}

            data: {"id":"c1","created":123,"model":"m","provider":"Baidu","choices":[{"index":1,"delta":{"content":"B"},"finish_reason":null}]}

            data: {"id":"c1","created":123,"model":"m","provider":"Baidu","choices":[{"index":0,"delta":{},"finish_reason":"stop"},{"index":1,"delta":{},"finish_reason":"stop"}]}

            data: {"id":"c1","created":123,"model":"m","provider":"Baidu","usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3}}

            data: [DONE]
        """.trimIndent()

        val out = Json.parseToJsonElement(rebuildFromChunks(sse)).jsonObject
        assertEquals("c1", out["id"]?.jsonPrimitive?.content)
        assertEquals("m", out["model"]?.jsonPrimitive?.content)
        assertEquals("Baidu", out["provider"]?.jsonPrimitive?.content)
        assertEquals("sf1", out["system_fingerprint"]?.jsonPrimitive?.content)
        assertEquals("chat.completion", out["object"]?.jsonPrimitive?.content)
        assertEquals(123, out["created"]?.jsonPrimitive?.content?.toLong())
        assertEquals(3, out["usage"]?.jsonObject?.get("total_tokens")?.jsonPrimitive?.content?.toInt())
        val choices = out["choices"]?.jsonArray ?: error("no choices")
        assertEquals(2, choices.size)
        assertEquals("Hello", choices[0].jsonObject["message"]?.jsonObject?.get("content")?.jsonPrimitive?.content)
        assertEquals("B", choices[1].jsonObject["message"]?.jsonObject?.get("content")?.jsonPrimitive?.content)
        assertEquals("stop", choices[0].jsonObject["finish_reason"]?.jsonPrimitive?.content)
    }

    @Test
    fun rebuildFromChunksReturnsUpstreamErrorVerbatim() {
        val sse = """
            data: {"id":"c1","error":{"message":"Unsupported field 'foo'","type":"invalid_request"},"choices":[]}

            data: [DONE]
        """.trimIndent()
        val out = Json.parseToJsonElement(rebuildFromChunks(sse)).jsonObject
        val err = out["error"]?.jsonObject ?: error("error block missing")
        assertEquals("Unsupported field 'foo'", err["message"]?.jsonPrimitive?.content)
    }

    @Test
    fun sha256MatchesKnownVectors() {
        assertEquals(
            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
            sha256Hex(""),
        )
        assertEquals(
            "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad",
            sha256Hex("abc"),
        )
        assertEquals(
            "248d6a61d20638b8e5c026930c3e6039a33ce45964ff2167f6ecedd419db06c1",
            sha256Hex("abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq"),
        )
    }

    @Test
    fun sessionPrefixHashesAreStableWhenHistoryGrows() {
        val first = Json.parseToJsonElement(
            """{"messages":[{"role":"system","content":"S"},{"role":"user","content":"hi"}]}""",
        ).jsonObject
        val second = Json.parseToJsonElement(
            """{"messages":[{"role":"system","content":"S"},{"role":"user","content":"hi"},{"role":"assistant","content":"yo"},{"role":"user","content":"again"}]}""",
        ).jsonObject

        val h1 = sessionPrefixHashes(first)
        val h2 = sessionPrefixHashes(second)

        // Цепочка начинается с первого user — system-преамбула не хэшируется.
        assertEquals(1, h1.size)
        assertEquals(3, h2.size)
        assertEquals(h1, h2.take(1))
        assertEquals(h1.last(), h2[0])
    }

    @Test
    fun sessionPrefixHashesIgnoreLeadingSystemSoSharedPromptDoesNotCollide() {
        val a = Json.parseToJsonElement(
            """{"messages":[{"role":"system","content":"SAME"},{"role":"user","content":"session A"}]}""",
        ).jsonObject
        val b = Json.parseToJsonElement(
            """{"messages":[{"role":"system","content":"SAME"},{"role":"user","content":"session B"}]}""",
        ).jsonObject
        // разные первые user-сообщения → разные хэши, несмотря на общий system
        assertFalse(sessionPrefixHashes(a) == sessionPrefixHashes(b))
    }

    @Test
    fun sessionPrefixHashesFallBackToWholeBodyWithoutMessages() {
        val body = Json.parseToJsonElement("""{"model":"m"}""").jsonObject
        assertEquals(listOf(sha256Hex(body.toString())), sessionPrefixHashes(body))
    }

    @Test
    fun sessionRegistryReusesIdByLongestCommonPrefix() {
        val reg = SessionRegistry()
        val first = reg.resolveLocked(listOf("H1"), 0)
        assertEquals("H1", first)

        // история выросла: [H1, H2, H3] — самый длинный известный префикс H1
        val next = reg.resolveLocked(listOf("H1", "H2", "H3"), 1)
        assertEquals("H1", next)

        // и дальше — id не меняется
        val deep = reg.resolveLocked(listOf("H1", "H2", "H3", "H4"), 2)
        assertEquals("H1", deep)
    }

    @Test
    fun sessionRegistryCreatesNewIdForDifferentHistory() {
        val reg = SessionRegistry()
        assertEquals("A1", reg.resolveLocked(listOf("A1"), 0))
        assertEquals("B1", reg.resolveLocked(listOf("B1"), 0))
        assertEquals("A1", reg.resolveLocked(listOf("A1", "A2"), 0))
    }

    @Test
    fun sessionRegistryEvictsLruWhenOverCapacity() {
        val reg = SessionRegistry(maxSessions = 2, ttlMillis = Long.MAX_VALUE)
        reg.resolveLocked(listOf("A1"), 0)
        reg.resolveLocked(listOf("B1"), 1)
        reg.resolveLocked(listOf("C1"), 2) // A вытеснена
        // a1 больше неизвестен → новая сессия с id = хэш всей истории A2
        assertEquals("A2", reg.resolveLocked(listOf("A1", "A2"), 3))
    }

    @Test
    fun sessionRegistryEvictsExpiredByTtl() {
        val reg = SessionRegistry(maxSessions = 100, ttlMillis = 1000)
        reg.resolveLocked(listOf("A1"), 0)
        reg.resolveLocked(listOf("B1"), 2000) // A протухла
        assertEquals("A2", reg.resolveLocked(listOf("A1", "A2"), 2001))
    }

    @Test
    fun parseConfigReadsSessionHeader() {
        val yaml = """
            providers:
              - id: p1
                url: "https://x.ru/api/v1"
                session_header: x-opencode-session
              - id: p2
                url: "https://y.ru/api/v1"
            models:
              - name: m1
                upstreams: []
        """.trimIndent()
        val cfg = parseConfig(Yaml.decodeYamlFromString(yaml))
        assertEquals("x-opencode-session", cfg.providers[0].session_header)
        assertEquals(null, cfg.providers[1].session_header)
    }

    @Test
    fun parseConfigReadsThinkTagsOnProviderAndUpstream() {
        val yaml = """
            providers:
              - id: p1
                url: "https://x.ru/api/v1"
                think_tags: split
              - id: p2
                url: "https://y.ru/api/v1"
            upstreams:
              - id: u1
                provider: p1
                model: real-1
                think_tags: strip
              - id: u2
                provider: p1
                model: real-2
            models:
              - name: m1
                upstreams: [u1, u2]
        """.trimIndent()

        val cfg = parseConfig(Yaml.decodeYamlFromString(yaml))

        assertEquals("split", cfg.providers[0].think_tags)
        assertEquals(null, cfg.providers[1].think_tags)
        assertEquals("strip", cfg.upstreams[0].think_tags)
        assertEquals(null, cfg.upstreams[1].think_tags)
    }

    @Test
    fun effectiveThinkTagsPrefersUpstreamThenProviderWithTolerantParse() {
        val upSplit = UpstreamConf("u1", "p", "m", think_tags = "split")
        val upNull = UpstreamConf("u2", "p", "m", think_tags = null)
        val prov = { tt: String? -> ProviderConf("p", "https://x", think_tags = tt) }

        // значение у апстрима — берётся оно, провайдер игнорируется
        assertEquals("split", effectiveThinkTags(upSplit, null))
        assertEquals("split", effectiveThinkTags(upSplit, prov("strip")))

        // апстрим null — берётся провайдерский
        assertEquals("split", effectiveThinkTags(upNull, prov("split")))
        assertEquals("strip", effectiveThinkTags(upNull, prov("strip")))
        assertEquals("split", effectiveThinkTags(upNull, prov("true")))
        assertEquals("off", effectiveThinkTags(upNull, prov("false")))
        assertEquals("off", effectiveThinkTags(upNull, prov("yes")))

        // нигде нет — "off"
        assertEquals("off", effectiveThinkTags(upNull, prov(null)))
        assertEquals("off", effectiveThinkTags(upNull, null))
    }

    @Test
    fun parseConfigReadsReasoningFieldAndEmptyOkOnProvider() {
        val yaml = """
            providers:
              - id: p1
                url: "https://x.ru/api/v1"
                reasoning_field: reasoning_content
                reasoning_empty_ok: true
              - id: p2
                url: "https://y.ru/api/v1"
                reasoning_empty_ok: false
            models:
              - name: m1
                upstreams: []
        """.trimIndent()
        val cfg = parseConfig(Yaml.decodeYamlFromString(yaml))
        assertEquals("reasoning_content", cfg.providers[0].reasoning_field)
        assertEquals(true, cfg.providers[0].reasoning_empty_ok)
        assertEquals(null, cfg.providers[1].reasoning_field)
        assertEquals(false, cfg.providers[1].reasoning_empty_ok)
    }

    @Test
    fun parseConfigDefaultsReasoningFieldsWhenAbsent() {
        val yaml = """
            providers:
              - id: p1
                url: "https://x.ru/api/v1"
            models:
              - name: m1
                upstreams: []
        """.trimIndent()
        val cfg = parseConfig(Yaml.decodeYamlFromString(yaml))
        assertEquals(null, cfg.providers[0].reasoning_field)
        assertEquals(false, cfg.providers[0].reasoning_empty_ok)
    }
}