package com.correx.apps.server.health

import io.ktor.client.HttpClient
import io.ktor.client.engine.mock.MockEngine
import io.ktor.client.engine.mock.respond
import io.ktor.http.HttpStatusCode
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Test

class LlamaServerHealthProbeTest {

    /** Replays a queued script of liveness readings; [error] is thrown when the queue runs dry on demand. */
    private class ScriptedClient(results: List<LivenessResult>, private val throwOnPing: Boolean = false) :
        LlamaLivenessClient {
        private val queue = ArrayDeque(results)
        override suspend fun ping(): LivenessResult {
            if (throwOnPing) error("ping boom")
            return queue.removeFirst()
        }
    }

    private fun reachable(tps: Double? = null) = LivenessResult(reachable = true, tokensPerSecond = tps)

    @Test
    fun `reachable with stable throughput stays healthy`(): Unit = runBlocking {
        val probe = LlamaServerHealthProbe(ScriptedClient(listOf(reachable(100.0), reachable(100.0), reachable(100.0))))
        assertEquals(HealthStatus.HEALTHY, probe.observe().status)
        assertEquals(HealthStatus.HEALTHY, probe.observe().status)
        val obs = probe.observe()
        assertEquals(HealthStatus.HEALTHY, obs.status)
        assertEquals("tokens_per_sec", obs.metric)
    }

    @Test
    fun `unreachable degrades on the liveness metric`(): Unit = runBlocking {
        val probe = LlamaServerHealthProbe(ScriptedClient(listOf(LivenessResult(reachable = false))))
        val obs = probe.observe()
        assertEquals(HealthStatus.DEGRADED, obs.status)
        assertEquals("liveness", obs.metric)
        assertEquals(0, obs.observedValue)
    }

    @Test
    fun `reachable without telemetry stays healthy on liveness`(): Unit = runBlocking {
        val probe = LlamaServerHealthProbe(ScriptedClient(listOf(reachable())))
        val obs = probe.observe()
        assertEquals(HealthStatus.HEALTHY, obs.status)
        assertEquals("liveness", obs.metric)
        assertEquals(1, obs.observedValue)
    }

    @Test
    fun `throughput collapse below half of peak degrades, then recovers`(): Unit = runBlocking {
        val probe = LlamaServerHealthProbe(
            ScriptedClient(listOf(reachable(100.0), reachable(100.0), reachable(30.0), reachable(100.0))),
        )
        assertEquals(HealthStatus.HEALTHY, probe.observe().status)
        assertEquals(HealthStatus.HEALTHY, probe.observe().status)

        val dropped = probe.observe()
        assertEquals(HealthStatus.DEGRADED, dropped.status)
        assertEquals("tokens_per_sec", dropped.metric)
        assertEquals(30, dropped.observedValue)

        val recovered = probe.observe()
        assertEquals(HealthStatus.HEALTHY, recovered.status)
        assertEquals("tokens_per_sec", recovered.metric)
    }

    @Test
    fun `ping that throws is treated as unreachable`(): Unit = runBlocking {
        val probe = LlamaServerHealthProbe(ScriptedClient(emptyList(), throwOnPing = true))
        val obs = probe.observe()
        assertEquals(HealthStatus.DEGRADED, obs.status)
        assertEquals("liveness", obs.metric)
    }

    private fun clientFor(status: HttpStatusCode) =
        HttpLlamaLivenessClient(HttpClient(MockEngine { respond(content = "", status = status) }), "http://llama")

    @Test
    fun `http client reports reachable on 200`(): Unit = runBlocking {
        assertEquals(true, clientFor(HttpStatusCode.OK).ping().reachable)
    }

    @Test
    fun `http client reports unreachable on 503`(): Unit = runBlocking {
        assertEquals(false, clientFor(HttpStatusCode.ServiceUnavailable).ping().reachable)
    }

    @Test
    fun `http client reports unreachable when transport throws`(): Unit = runBlocking {
        val client = HttpLlamaLivenessClient(
            HttpClient(MockEngine { error("connection refused") }),
            "http://llama",
        )
        assertEquals(false, client.ping().reachable)
    }
}
