package com.correx.infrastructure.inference.commons

import org.slf4j.LoggerFactory
import java.io.File

/**
 * Whole-device GPU usage at an instant. All fields are device-wide (not per-process) —
 * per-model VRAM attribution is an explicit non-goal.
 */
data class GpuStatus(
    val memoryUsedMb: Long,
    val memoryTotalMb: Long,
    val utilizationPct: Int,
)

/**
 * A point-in-time reading of local resource usage. All fields are nullable so that "unavailable"
 * is representable without exceptions: [gpu] is null when no GPU probe is present (non-NVIDIA/AMD
 * box or the tool is absent), [processRssBytes] is null when no correx-managed model is resident,
 * and [systemRamUsedBytes]/[systemRamTotalBytes] are null when `/proc/meminfo` can't be read.
 *
 * System RAM is **owner-agnostic** — it reflects the whole host regardless of who launched the
 * model, so the gauge renders even on the static (external `llama-server`) path.
 */
data class ResourceSnapshot(
    val gpu: GpuStatus?,
    val processRssBytes: Long?,
    val systemRamUsedBytes: Long? = null,
    val systemRamTotalBytes: Long? = null,
) {
    companion object {
        val UNAVAILABLE = ResourceSnapshot(gpu = null, processRssBytes = null)
    }
}

/**
 * Vendor-agnostic live gauge of GPU/process resource usage. This is a **live gauge, never recorded
 * as events** — it feeds no core decision, so invariant #9 does not require event-sourcing it, and
 * replay (invariant #8) never reads it. Implementations must degrade gracefully (return
 * [ResourceSnapshot.UNAVAILABLE]-shaped data) rather than throw.
 */
fun interface ResourceProbe {
    fun probe(): ResourceSnapshot
}

/** Always-unavailable probe — used on non-GPU boxes when even system RAM can't be read. */
object UnavailableProbe : ResourceProbe {
    override fun probe(): ResourceSnapshot = ResourceSnapshot.UNAVAILABLE
}

/**
 * Owner-agnostic composite probe: delegates GPU/process-RSS to a vendor [gpuProbe] (NVIDIA/AMD or
 * [UnavailableProbe]) and always overlays whole-host system RAM read from `/proc/meminfo`. This is
 * what makes the gauge render on BOTH the managed and static (external `llama-server`) paths — the
 * system-RAM read needs no managed pid. The RAM read is injectable for testing.
 */
class SystemResourceProbe(
    private val gpuProbe: ResourceProbe = UnavailableProbe,
    private val systemRamReader: () -> Pair<Long, Long>? = SystemRam::readUsedTotalBytes,
) : ResourceProbe {
    override fun probe(): ResourceSnapshot {
        val base = gpuProbe.probe()
        val ram = systemRamReader()
        return base.copy(
            systemRamUsedBytes = ram?.first,
            systemRamTotalBytes = ram?.second,
        )
    }
}

/** Reads whole-host RAM from `/proc/meminfo`. Fails soft to null so a non-Linux host never throws. */
object SystemRam {
    private val log = LoggerFactory.getLogger("SystemRam")
    private const val BYTES_PER_KB = 1024L

    /** Returns (usedBytes, totalBytes) where used = MemTotal − MemAvailable, or null if unreadable. */
    fun readUsedTotalBytes(): Pair<Long, Long>? = runCatching {
        val meminfo = File("/proc/meminfo")
        if (!meminfo.exists()) return null
        val lines = meminfo.readLines()
        val total = lines.kbValue("MemTotal:") ?: return null
        val available = lines.kbValue("MemAvailable:") ?: return null
        val used = (total - available).coerceAtLeast(0)
        (used * BYTES_PER_KB) to (total * BYTES_PER_KB)
    }.onFailure { log.debug("/proc/meminfo read failed: {}", it.message) }.getOrNull()

    // Parses a "Key:   12345 kB" line to its kB value.
    private fun List<String>.kbValue(key: String): Long? =
        firstOrNull { it.startsWith(key) }
            ?.removePrefix(key)?.trim()?.substringBefore(' ')?.toLongOrNull()
}
