package com.correx.apps.server.health

import com.correx.core.events.events.HealthSubject
import com.correx.core.events.events.InferenceCompletedEvent
import com.correx.core.events.stores.EventStore
import io.ktor.client.HttpClient
import io.ktor.client.request.get
import io.ktor.http.isSuccess
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import kotlinx.coroutines.withTimeout

private const val MS_PER_SECOND = 1000.0
private const val TPS_WINDOW = 10

/**
 * Health probe for the local llama-server (observability-spec §4: LLAMA_SERVER subject).
 *
 * **Liveness** — every tick issues a bounded `GET [llamaBaseUrl]/health` via [httpClient].
 * A 200-family response is live; any non-2xx, connection error, or [livenessTimeoutMs]
 * expiry is DEGRADED. The [httpClient] is injected so tests can supply a [MockEngine].
 *
 * **Degradation (tokens/sec)** — on each tick the probe scans new [InferenceCompletedEvent]s
 * from the event store (events appended since the previous tick, tracked by
 * [lastScannedGlobalSequence]), accumulates the per-call tokens/sec into a sliding window
 * of size [TPS_WINDOW], and reports DEGRADED when the window average falls below
 * [tpsWarnBelow]. A threshold of 0.0 disables the degradation dimension; a window that
 * has not yet collected any samples reports HEALTHY (no baseline).
 *
 * The probe only observes and returns a [HealthObservation]; it never emits events.
 * [HealthMonitor] owns edge detection and event emission (Invariants #8/#9).
 */
class LlamaServerHealthProbe(
    private val llamaBaseUrl: String,
    private val httpClient: HttpClient,
    private val eventStore: EventStore,
    private val livenessTimeoutMs: Long = 5_000,
    private val tpsWarnBelow: Double = 0.0,
) : HealthProbe {

    override val subject: HealthSubject = HealthSubject.LLAMA_SERVER

    /**
     * Monotonic global sequence up to which we have already scanned for inference events.
     * Assumes single-coroutine sequential invocation: [HealthMonitor] calls probes sequentially
     * within one coroutine tick, so no concurrent mutation of this field or [tpsWindow] occurs.
     */
    private var lastScannedGlobalSequence: Long = -1L

    /**
     * Sliding window of recent per-call tokens/sec values (newest at the end).
     * See [lastScannedGlobalSequence] for the single-coroutine assumption.
     */
    private val tpsWindow: ArrayDeque<Double> = ArrayDeque(TPS_WINDOW)

    override suspend fun observe(): HealthObservation {
        val livenessObservation = checkLiveness()
        if (livenessObservation.status == HealthStatus.DEGRADED) return livenessObservation
        return checkThroughput() ?: livenessObservation
    }

    private suspend fun checkLiveness(): HealthObservation {
        val result = runCatching {
            withTimeout(livenessTimeoutMs) {
                httpClient.get("$llamaBaseUrl/health")
            }
        }
        val response = result.getOrNull()
        return when {
            result.isFailure -> degraded(
                metric = "liveness",
                value = 0L,
                threshold = 1L,
                detail = "GET $llamaBaseUrl/health failed: ${result.exceptionOrNull()?.message}",
            )
            response != null && !response.status.isSuccess() -> degraded(
                metric = "liveness",
                value = response.status.value.toLong(),
                threshold = 200L,
                detail = "GET $llamaBaseUrl/health returned HTTP ${response.status.value}",
            )
            else -> healthy(
                metric = "liveness",
                value = 1L,
                threshold = 1L,
                detail = "llama-server live at $llamaBaseUrl",
            )
        }
    }

    private suspend fun checkThroughput(): HealthObservation? {
        if (tpsWarnBelow <= 0.0) {
            // Dimension disabled: still advance the cursor so a later live-enable doesn't rescan
            // the entire history from sequence 0.
            lastScannedGlobalSequence = withContext(Dispatchers.IO) { eventStore.lastGlobalSequence() }
            return null
        }
        withContext(Dispatchers.IO) { scanNewInferenceEvents() }
        if (tpsWindow.isEmpty()) return null
        val avgTps = tpsWindow.average()
        val avgTpsLong = (avgTps * TPS_SCALE).toLong()
        val thresholdLong = (tpsWarnBelow * TPS_SCALE).toLong()
        return if (avgTps < tpsWarnBelow) {
            degraded(
                metric = "tokens_per_second",
                value = avgTpsLong,
                threshold = thresholdLong,
                detail = "avg tokens/sec %.2f below warn threshold %.2f (window=%d)".format(
                    avgTps, tpsWarnBelow, tpsWindow.size,
                ),
            )
        } else {
            healthy(
                metric = "tokens_per_second",
                value = avgTpsLong,
                threshold = thresholdLong,
                detail = "avg tokens/sec %.2f (window=%d)".format(avgTps, tpsWindow.size),
            )
        }
    }

    /**
     * Scans [EventStore.allEvents] for [InferenceCompletedEvent]s with a global sequence
     * number beyond [lastScannedGlobalSequence], appending each call's tokens/sec to
     * [tpsWindow] (capped at [TPS_WINDOW]). Updates [lastScannedGlobalSequence] so the
     * next tick only scans new events — O(new events), not O(all events).
     */
    private fun scanNewInferenceEvents() {
        var maxSeen = lastScannedGlobalSequence
        eventStore.allEvents()
            .filter { it.sequence > lastScannedGlobalSequence }
            .forEach { event ->
                if (event.sequence > maxSeen) maxSeen = event.sequence
                val payload = event.payload as? InferenceCompletedEvent ?: return@forEach
                // Skip events with no completion tokens (e.g. stop-token-only responses) to
                // avoid dragging the windowed average to zero on a healthy model.
                if (payload.tokensUsed.completionTokens <= 0) return@forEach
                val tps = if (payload.latencyMs > 0L) {
                    payload.tokensUsed.completionTokens.toDouble() / (payload.latencyMs / MS_PER_SECOND)
                } else {
                    0.0
                }
                if (tpsWindow.size >= TPS_WINDOW) tpsWindow.removeFirst()
                tpsWindow.addLast(tps)
            }
        lastScannedGlobalSequence = maxSeen
    }

    private fun degraded(metric: String, value: Long, threshold: Long, detail: String) =
        HealthObservation(subject, HealthStatus.DEGRADED, metric, value, threshold, detail)

    private fun healthy(metric: String, value: Long, threshold: Long, detail: String) =
        HealthObservation(subject, HealthStatus.HEALTHY, metric, value, threshold, detail)

    companion object {
        /** Scale factor: store tps×100 as Long to fit [HealthObservation.observedValue]. */
        private const val TPS_SCALE = 100L
    }
}
