package com.correx.infrastructure.inference.commons

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

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

/**
 * [ResourceProbe] backed by `rocm-smi --showmeminfo vram --showuse --csv` for AMD GPU stats and
 * `/proc/<pid>/status` for the managed model process RSS. Both reads fail soft to null, so a
 * missing `rocm-smi`, a non-AMD host, or a dead pid never throws.
 *
 * `rocm-smi` prefixes its CSV with warning/blank lines, and the column order/names vary between
 * ROCm versions — so the parser locates the header by its column names rather than by position.
 * VRAM figures are reported in bytes and converted to MiB to match [GpuStatus].
 *
 * The command runner and RSS read are injectable so the parsing is unit-testable without a GPU.
 *
 * @param pidSupplier returns the managed `llama-server` pid, or null when no model is resident.
 * @param rocmSmiRunner returns raw `rocm-smi --csv` stdout, or null when unavailable.
 * @param rssReader returns the RSS in bytes for a pid, or null when unreadable.
 */
class AmdResourceProbe(
    private val pidSupplier: () -> Long?,
    private val rocmSmiRunner: () -> String? = ::runRocmSmi,
    private val rssReader: (Long) -> Long? = ::readRssBytes,
) : ResourceProbe {

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

    companion object {
        private val QUERY = listOf(
            "rocm-smi", "--showmeminfo", "vram", "--showuse", "--csv",
        )
        private const val BYTES_PER_MB = 1024L * 1024L
        private const val BYTES_PER_KB = 1024L

        /**
         * Parses `rocm-smi --csv` output. Skips the warning/blank preamble, finds the header by its
         * column names (so column order and version differences don't matter), reads the first card
         * row, and converts the byte VRAM figures to MiB. Returns null on any shape mismatch so the
         * caller surfaces "unavailable" rather than partial data.
         *
         * Example input (AMD GPU in low-power state warning elided):
         * ```
         * device,GPU use (%),VRAM Total Memory (B),VRAM Total Used Memory (B)
         * card0,2,17163091968,12186726400
         * ```
         */
        fun parseGpu(csv: String): GpuStatus? = runCatching {
            val lines = csv.lineSequence().map { it.trim() }.filter { it.isNotEmpty() }.toList()
            val headerIdx = lines.indexOfFirst {
                it.startsWith("device,") && it.contains("VRAM Total", ignoreCase = true)
            }
            require(headerIdx >= 0 && headerIdx + 1 < lines.size)

            val header = lines[headerIdx].split(',').map { it.trim() }
            val row = lines[headerIdx + 1].split(',').map { it.trim() }

            val usedCol = header.indexOfFirst { it.contains("VRAM", true) && it.contains("Used", true) }
            val totalCol = header.indexOfFirst { it.contains("VRAM", true) && it.contains("Total Memory", true) }
            val useCol = header.indexOfFirst { it.contains("use (%)", true) }
            require(usedCol >= 0 && totalCol >= 0)

            val util = useCol.takeIf { it >= 0 }?.let { row.getOrNull(it)?.toIntOrNull() } ?: 0
            GpuStatus(
                memoryUsedMb = row[usedCol].toLong() / BYTES_PER_MB,
                memoryTotalMb = row[totalCol].toLong() / BYTES_PER_MB,
                utilizationPct = util,
            )
        }.getOrNull()

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

        private fun runRocmSmi(): 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("rocm-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()
    }
}
