package com.correx.apps.server.health

import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.InferenceCompletedEvent
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.EventId
import com.correx.core.events.types.SessionId
import com.correx.core.inference.TokenUsage
import com.correx.core.utils.TypeId
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.flow.Flow
import kotlinx.coroutines.flow.emptyFlow
import kotlinx.coroutines.runBlocking
import kotlinx.datetime.Instant
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test

class LlamaServerHealthProbeTest {

    private val sessionId: SessionId = TypeId("s-1")
    private val stageId = TypeId("stage-1")

    // ---------------------------------------------------------------------------
    // Helpers
    // ---------------------------------------------------------------------------

    private fun httpClient(status: HttpStatusCode): HttpClient =
        HttpClient(MockEngine { respond("", status) })

    private fun httpClientThrowing(): HttpClient =
        HttpClient(MockEngine { error("connection refused") })

    private fun inference(seq: Long, completionTokens: Int, latencyMs: Long): StoredEvent =
        StoredEvent(
            metadata = EventMetadata(
                eventId = EventId("e$seq"),
                sessionId = sessionId,
                timestamp = Instant.parse("2026-01-01T00:00:00Z"),
                schemaVersion = 1,
                causationId = null,
                correlationId = null,
            ),
            sequence = seq,
            sessionSequence = seq,
            payload = InferenceCompletedEvent(
                requestId = TypeId("req-$seq"),
                sessionId = sessionId,
                stageId = stageId,
                providerId = TypeId("lfm"),
                tokensUsed = TokenUsage(promptTokens = 10, completionTokens = completionTokens),
                latencyMs = latencyMs,
                responseArtifactId = TypeId("art-$seq"),
            ),
        )

    /** Minimal EventStore backed by a mutable list; allEvents() returns all stored events. */
    private class ListEventStore(private val events: MutableList<StoredEvent> = mutableListOf()) : EventStore {
        override suspend fun append(event: NewEvent): StoredEvent = error("unused")
        override suspend fun appendAll(events: List<NewEvent>): List<StoredEvent> = error("unused")
        override fun read(sessionId: SessionId): List<StoredEvent> =
            events.filter { it.metadata.sessionId == sessionId }
        override fun readFrom(sessionId: SessionId, fromSequence: Long): List<StoredEvent> =
            events.filter { it.metadata.sessionId == sessionId && it.sequence >= fromSequence }
        override fun lastSequence(sessionId: SessionId): Long? =
            events.filter { it.metadata.sessionId == sessionId }.lastOrNull()?.sequence
        override fun subscribe(sessionId: SessionId): Flow<StoredEvent> = emptyFlow()
        override fun subscribeAll(): Flow<StoredEvent> = emptyFlow()
        override suspend fun lastGlobalSequence(): Long = events.lastOrNull()?.sequence ?: 0L
        override fun allEvents(): Sequence<StoredEvent> = events.asSequence()
        override fun allSessionIds(): Set<SessionId> = events.map { it.metadata.sessionId }.toSet()

        fun add(event: StoredEvent) { events.add(event) }
    }

    private fun probe(
        status: HttpStatusCode = HttpStatusCode.OK,
        store: ListEventStore = ListEventStore(),
        tpsWarnBelow: Double = 0.0,
    ) = LlamaServerHealthProbe(
        llamaBaseUrl = "http://localhost:11434",
        httpClient = httpClient(status),
        eventStore = store,
        livenessTimeoutMs = 5_000,
        tpsWarnBelow = tpsWarnBelow,
    )

    // ---------------------------------------------------------------------------
    // Liveness tests
    // ---------------------------------------------------------------------------

    @Test
    fun `healthy when llama-server returns 200`(): Unit = runBlocking {
        val obs = probe(status = HttpStatusCode.OK).observe()
        assertEquals(HealthStatus.HEALTHY, obs.status)
        assertEquals("liveness", obs.metric)
        assertEquals(1L, obs.observedValue)
    }

    @Test
    fun `degraded when llama-server returns non-200`(): Unit = runBlocking {
        val obs = probe(status = HttpStatusCode.ServiceUnavailable).observe()
        assertEquals(HealthStatus.DEGRADED, obs.status)
        assertEquals("liveness", obs.metric)
        assertTrue(obs.detail.contains("503"), "expected 503 in detail: ${obs.detail}")
    }

    @Test
    fun `degraded on connection error`(): Unit = runBlocking {
        val p = LlamaServerHealthProbe(
            llamaBaseUrl = "http://localhost:11434",
            httpClient = httpClientThrowing(),
            eventStore = ListEventStore(),
            livenessTimeoutMs = 5_000,
            tpsWarnBelow = 0.0,
        )
        val obs = p.observe()
        assertEquals(HealthStatus.DEGRADED, obs.status)
        assertEquals("liveness", obs.metric)
    }

    // ---------------------------------------------------------------------------
    // Tokens/sec degradation tests
    // ---------------------------------------------------------------------------

    @Test
    fun `healthy when tps above threshold`(): Unit = runBlocking {
        // 100 tokens in 1000 ms = 100 tps, threshold is 50 tps
        val store = ListEventStore()
        store.add(inference(seq = 1, completionTokens = 100, latencyMs = 1_000))
        val obs = probe(store = store, tpsWarnBelow = 50.0).observe()
        assertEquals(HealthStatus.HEALTHY, obs.status)
        assertEquals("tokens_per_second", obs.metric)
    }

    @Test
    fun `degraded when tps below threshold`(): Unit = runBlocking {
        // 10 tokens in 1000 ms = 10 tps, threshold is 50 tps
        val store = ListEventStore()
        store.add(inference(seq = 1, completionTokens = 10, latencyMs = 1_000))
        val obs = probe(store = store, tpsWarnBelow = 50.0).observe()
        assertEquals(HealthStatus.DEGRADED, obs.status)
        assertEquals("tokens_per_second", obs.metric)
        assertTrue(obs.detail.contains("below"), "detail should mention 'below': ${obs.detail}")
    }

    @Test
    fun `healthy when no inference events and tpsWarnBelow configured`(): Unit = runBlocking {
        // No baseline — can't degrade without samples
        val obs = probe(tpsWarnBelow = 50.0).observe()
        assertEquals(HealthStatus.HEALTHY, obs.status)
        // No inference events means the liveness path dominates
        assertEquals("liveness", obs.metric)
    }

    @Test
    fun `tpsWarnBelow zero disables degradation check`(): Unit = runBlocking {
        val store = ListEventStore()
        store.add(inference(seq = 1, completionTokens = 1, latencyMs = 100_000)) // very slow
        val obs = probe(store = store, tpsWarnBelow = 0.0).observe()
        // Degradation disabled; liveness is healthy -> overall HEALTHY
        assertEquals(HealthStatus.HEALTHY, obs.status)
        assertEquals("liveness", obs.metric)
    }

    @Test
    fun `only scans new events on subsequent ticks`(): Unit = runBlocking {
        val store = ListEventStore()
        store.add(inference(seq = 1, completionTokens = 100, latencyMs = 1_000)) // 100 tps

        val p = LlamaServerHealthProbe(
            llamaBaseUrl = "http://localhost:11434",
            httpClient = httpClient(HttpStatusCode.OK),
            eventStore = store,
            livenessTimeoutMs = 5_000,
            tpsWarnBelow = 50.0,
        )

        // First tick: sees inference at seq=1, 100 tps -> HEALTHY
        assertEquals(HealthStatus.HEALTHY, p.observe().status)

        // Add a very slow inference at seq=2
        store.add(inference(seq = 2, completionTokens = 1, latencyMs = 10_000)) // 0.1 tps

        // Second tick: sees only seq=2 (new), window = [100, 0.1], avg ~50.05 -> still HEALTHY
        val obs2 = p.observe()
        // avg(100, 0.1) = 50.05 which is above 50 threshold -> HEALTHY
        assertEquals(HealthStatus.HEALTHY, obs2.status)

        // Add another very slow inference
        store.add(inference(seq = 3, completionTokens = 1, latencyMs = 10_000)) // 0.1 tps

        // Third tick: window = [100, 0.1, 0.1], avg ~33.4 -> DEGRADED
        val obs3 = p.observe()
        assertEquals(HealthStatus.DEGRADED, obs3.status)
        assertEquals("tokens_per_second", obs3.metric)
    }
}
