package com.correx.testing.contracts.model.inference

import com.correx.infrastructure.inference.commons.AmdResourceProbe
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertNull
import org.junit.jupiter.api.Test

class AmdResourceProbeTest {

    // Real `rocm-smi --showmeminfo vram --showuse --csv` output (ROCm 4.0.0 / lib 7.8.0),
    // including the low-power warning + blank line that precede the CSV.
    private val realCsv = """
        WARNING: AMD GPU device(s) is/are in a low-power state. Check power control/runtime_status

        device,GPU use (%),VRAM Total Memory (B),VRAM Total Used Memory (B)
        card0,2,17163091968,12186726400
    """.trimIndent()

    @Test
    fun `parses real rocm-smi csv, converting bytes to MiB`() {
        val gpu = AmdResourceProbe.parseGpu(realCsv)
        assertEquals(16368L, gpu?.memoryTotalMb) // 17163091968 B / 1 MiB
        assertEquals(11622L, gpu?.memoryUsedMb) // 12186726400 B / 1 MiB
        assertEquals(2, gpu?.utilizationPct)
    }

    @Test
    fun `tolerates reordered columns by matching header names`() {
        val reordered = """
            device,VRAM Total Used Memory (B),VRAM Total Memory (B),GPU use (%)
            card0,12186726400,17163091968,2
        """.trimIndent()
        val gpu = AmdResourceProbe.parseGpu(reordered)
        assertEquals(16368L, gpu?.memoryTotalMb)
        assertEquals(11622L, gpu?.memoryUsedMb)
        assertEquals(2, gpu?.utilizationPct)
    }

    @Test
    fun `returns null when header or data is missing`() {
        assertNull(AmdResourceProbe.parseGpu(""))
        assertNull(AmdResourceProbe.parseGpu("WARNING: something\n\n"))
        assertNull(AmdResourceProbe.parseGpu("device,GPU use (%),VRAM Total Memory (B),VRAM Total Used Memory (B)"))
    }

    @Test
    fun `probe combines gpu csv and rss reader`() {
        val probe = AmdResourceProbe(
            pidSupplier = { 4242L },
            rocmSmiRunner = { realCsv },
            rssReader = { pid -> if (pid == 4242L) 2_097_152L else null },
        )
        val snapshot = probe.probe()
        assertEquals(11622L, snapshot.gpu?.memoryUsedMb)
        assertEquals(2_097_152L, snapshot.processRssBytes)
    }

    @Test
    fun `probe reports gpu unavailable when rocm-smi yields null`() {
        val probe = AmdResourceProbe(
            pidSupplier = { null },
            rocmSmiRunner = { null },
            rssReader = { error("must not be called when pid is null") },
        )
        val snapshot = probe.probe()
        assertNull(snapshot.gpu)
        assertNull(snapshot.processRssBytes)
    }
}
