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

package pw.binom.llmproxy

import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import kotlinx.coroutines.test.TestScope
import kotlinx.coroutines.test.runTest
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.jsonArray
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.jsonPrimitive
import kotlin.test.Test
import kotlin.test.assertEquals
import kotlin.test.assertNull
import kotlin.time.Duration.Companion.milliseconds
import kotlin.time.Duration.Companion.minutes
import kotlin.time.Duration.Companion.seconds

class HealthProberTest {

    private val json = Json { ignoreUnknownKeys = true }

    @Test
    fun effectiveProbeIntervalPrefersUpstreamThenProviderThenDefault() {
        val provider = ProviderConf("p", "https://x", "", null, null, probe_interval = 2.seconds)
        val upstream = UpstreamConf("u", "p", "m", null, null, null, null, null, probe_interval = 5.milliseconds)
        val upstreamNoOwn = UpstreamConf("u2", "p", "m")
        assertEquals(5.milliseconds, effectiveProbeInterval(upstream, provider))
        assertEquals(2.seconds, effectiveProbeInterval(upstreamNoOwn, provider))
        assertEquals(DEFAULT_PROBE_INTERVAL, effectiveProbeInterval(upstreamNoOwn, null))
    }

    @Test
    fun effectiveProbeIntervalZeroDisablesProber() {
        val up = UpstreamConf("u", "p", "m", null, null, null, null, null, probe_interval = 0.milliseconds)
        assertNull(effectiveProbeInterval(up, null))

        val providerZero = ProviderConf("p", "https://x", "", null, null, probe_interval = 0.milliseconds)
        assertNull(effectiveProbeInterval(UpstreamConf("u2", "p", "m"), providerZero))
    }

    @Test
    fun buildProbeBodyIsMinimalChatWithModelReplacedAndPatchesApplied() {
        val provider = ProviderConf(
            "p", "https://x", "", null,
            Json.parseToJsonElement("""{"provider":{"allow_fallbacks":false}}""").jsonObject,
            probe_interval = 1.seconds,
        )
        val up = UpstreamConf(
            "u", "p", "real-1", null,
            Json.parseToJsonElement("""{"temperature":0.5}""").jsonObject,
            null, null, null, 1.seconds,
        )

        val body = buildProbeBody(up, provider)

        assertEquals("real-1", body["model"]?.jsonPrimitive?.content)
        assertEquals("2+2=?", body["messages"]?.jsonArray?.first()?.jsonObject?.get("content")?.jsonPrimitive?.content)
        assertEquals("true", body["stream"]?.jsonPrimitive?.content)
        assertEquals("4", body["max_tokens"]?.jsonPrimitive?.content)
        assertEquals("false", body["provider"]?.jsonObject?.get("allow_fallbacks")?.jsonPrimitive?.content)
        assertEquals("0.5", body["temperature"]?.jsonPrimitive?.content)
    }

    @Test
    fun buildProbeBodyIgnoresModelPatch() {
        val up = UpstreamConf("u", "p", "real-1")
        val provider = ProviderConf("p", "https://x")
        val body = buildProbeBody(up, provider)
        assertEquals(
            """{"model":"real-1","messages":[{"role":"user","content":"2+2=?"}],"max_tokens":4,"stream":true}""",
            body.toString(),
        )
    }

    @Test
    fun probeLoopRecoversUpstreamWhenPingSucceeds() = runTest {
        val provider = ProviderConf("p", "https://x", "", null, null, backoff = 5.minutes)
        val up = UpstreamConf("u", "p", "m", null, null, null, backoff = 5.minutes, probe_interval = 5.milliseconds)
        val backoff = BackoffRegistry(mapOf("p" to provider), mapOf("u" to up))
        backoff.recordFailure(up)
        assertEquals(true, backoff.isCooling(up))

        val metrics = MetricsRegistry()
        val prober = HealthProber(
            providersById = mapOf("p" to provider),
            upstreams = listOf(up),
            http = createHttpClient(),
            backoff = backoff,
            metrics = metrics,
            scope = this,
        )
        var pings = 0
        val job = launch {
            prober.probeLoop(up, 5.milliseconds) { pinged ->
                pings++
                assertEquals("u", pinged.id)
                true
            }
        }
        runUntil { pings >= 1 }
        // успех уже зафиксирован (recordSuccess) — даём циклу несколько
        // итераций и убеждаемся, что повторных пингов не было
        delay(10.milliseconds)
        job.cancel()

        assertEquals(1, pings)
        assertEquals(false, backoff.isCooling(up))
    }

    @Test
    fun probeLoopKeepsCoolingWhenPingFails() = runTest {
        val provider = ProviderConf("p", "https://x", "", null, null, backoff = 5.minutes)
        val up = UpstreamConf("u", "p", "m", null, null, null, backoff = 5.minutes, probe_interval = 5.milliseconds)
        val backoff = BackoffRegistry(mapOf("p" to provider), mapOf("u" to up))
        backoff.recordFailure(up)

        val metrics = MetricsRegistry()
        val prober = HealthProber(
            providersById = mapOf("p" to provider),
            upstreams = listOf(up),
            http = createHttpClient(),
            backoff = backoff,
            metrics = metrics,
            scope = this,
        )
        var pings = 0
        val job = launch {
            prober.probeLoop(up, 5.milliseconds) { _ ->
                pings++
                false
            }
        }
        runUntil { pings >= 2 }
        job.cancel()

        assertEquals(true, backoff.isCooling(up))
    }

    private suspend fun TestScope.runUntil(condition: () -> Boolean) {
        var i = 0
        while (!condition() && i < 10_000) {
            delay(1)
            i++
        }
        if (!condition()) throw AssertionError("condition not met in time")
    }

}