package com.correx.infrastructure.inference.commons

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

private val log = LoggerFactory.getLogger("NvidiaResourceProbe")

/**
 * [ResourceProbe] backed by `nvidia-smi` for GPU stats and `/proc/<pid>/status` for the managed
 * model process RSS. Both reads are isolated and fail soft to null, so a missing `nvidia-smi`,
 * a non-NVIDIA host, or a dead pid never throws.
 *
 * The `nvidia-smi` invocation and RSS read are injectable so the parsing is unit-testable without
 * a GPU or a live process.
 *
 * @param pidSupplier returns the managed `llama-server` pid, or null when no model is resident.
 * @param nvidiaSmiRunner returns raw `nvidia-smi` CSV stdout, or null when unavailable.
 * @param rssReader returns the RSS in bytes for a pid, or null when unreadable.
 */
class NvidiaResourceProbe(
    private val pidSupplier: () -> Long?,
    private val nvidiaSmiRunner: () -> String? = ::runNvidiaSmi,
    private val rssReader: (Long) -> Long? = ::readRssBytes,
) : ResourceProbe {

    override fun probe(): ResourceSnapshot {
        val gpu = nvidiaSmiRunner()?.let { parseGpu(it) }
        val rss = pidSupplier()?.let { rssReader(it) }
        return ResourceSnapshot(gpu = gpu, processRssBytes = rss)
    }

    companion object {
        private const val EXPECTED_FIELDS = 3
        private val QUERY = listOf(
            "nvidia-smi",
            "--query-gpu=memory.used,memory.total,utilization.gpu",
            "--format=csv,noheader,nounits",
        )

        /**
         * Parses the first GPU row of `nvidia-smi --format=csv,noheader,nounits` output
         * (`memory.used, memory.total, utilization.gpu`, all in MiB / percent). Returns null on any
         * shape mismatch so the caller surfaces "unavailable" instead of partial data.
         */
        fun parseGpu(csv: String): GpuStatus? {
            val row = csv.lineSequence().map { it.trim() }.firstOrNull { it.isNotEmpty() } ?: return null
            val fields = row.split(',').map { it.trim() }
            if (fields.size < EXPECTED_FIELDS) return null
            val used = fields[0].toLongOrNull() ?: return null
            val total = fields[1].toLongOrNull() ?: return null
            val util = fields[2].toIntOrNull() ?: return null
            return GpuStatus(memoryUsedMb = used, memoryTotalMb = total, utilizationPct = util)
        }

        /** Returns true if `nvidia-smi -L` runs successfully — used to pick the probe at boot. */
        fun isAvailable(): Boolean = runCatching {
            val proc = ProcessBuilder("nvidia-smi", "-L")
                .redirectErrorStream(true)
                .start()
            proc.waitFor() == 0
        }.getOrElse { false }

        private fun runNvidiaSmi(): String? = runCatching {
            val proc = ProcessBuilder(QUERY).redirectErrorStream(true).start()
            val out = proc.inputStream.bufferedReader().readText()
            if (proc.waitFor() == 0) out else null
        }.onFailure { log.debug("nvidia-smi query failed: {}", it.message) }.getOrNull()

        private fun readRssBytes(pid: Long): Long? = runCatching {
            val statusFile = File("/proc/$pid/status")
            if (!statusFile.exists()) return null
            val vmRss = statusFile.readLines().firstOrNull { it.startsWith("VmRSS:") } ?: return null
            // Format: "VmRSS:\t   12345 kB"
            val kb = vmRss.removePrefix("VmRSS:").trim().substringBefore(' ').toLongOrNull() ?: return null
            kb * BYTES_PER_KB
        }.onFailure { log.debug("RSS read failed for pid {}: {}", pid, it.message) }.getOrNull()

        private const val BYTES_PER_KB = 1024L
    }
}
