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

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

class NvidiaResourceProbeTest {

    @Test
    fun `parses a single GPU csv row`() {
        val gpu = NvidiaResourceProbe.parseGpu("8192, 16384, 37")
        assertEquals(8192L, gpu?.memoryUsedMb)
        assertEquals(16384L, gpu?.memoryTotalMb)
        assertEquals(37, gpu?.utilizationPct)
    }

    @Test
    fun `parses the first row when multiple GPUs are present`() {
        val gpu = NvidiaResourceProbe.parseGpu("8192, 16384, 37\n1024, 16384, 5")
        assertEquals(8192L, gpu?.memoryUsedMb)
        assertEquals(37, gpu?.utilizationPct)
    }

    @Test
    fun `returns null on malformed csv`() {
        assertNull(NvidiaResourceProbe.parseGpu("not,enough"))
        assertNull(NvidiaResourceProbe.parseGpu("abc, def, ghi"))
        assertNull(NvidiaResourceProbe.parseGpu(""))
    }

    @Test
    fun `probe combines gpu csv and rss reader`() {
        val probe = NvidiaResourceProbe(
            pidSupplier = { 4242L },
            nvidiaSmiRunner = { "8192, 16384, 37" },
            rssReader = { pid -> if (pid == 4242L) 1_048_576L else null },
        )
        val snapshot = probe.probe()
        assertEquals(8192L, snapshot.gpu?.memoryUsedMb)
        assertEquals(1_048_576L, snapshot.processRssBytes)
    }

    @Test
    fun `probe reports gpu unavailable when nvidia-smi yields null`() {
        val probe = NvidiaResourceProbe(
            pidSupplier = { 4242L },
            nvidiaSmiRunner = { null },
            rssReader = { 1_048_576L },
        )
        val snapshot = probe.probe()
        assertNull(snapshot.gpu)
        assertEquals(1_048_576L, snapshot.processRssBytes)
    }

    @Test
    fun `probe reports rss null when no pid is resident`() {
        val probe = NvidiaResourceProbe(
            pidSupplier = { null },
            nvidiaSmiRunner = { "8192, 16384, 37" },
            rssReader = { error("must not be called when pid is null") },
        )
        val snapshot = probe.probe()
        assertEquals(8192L, snapshot.gpu?.memoryUsedMb)
        assertNull(snapshot.processRssBytes)
    }

    @Test
    fun `unavailable probe is fully unavailable`() {
        val snapshot = UnavailableProbe.probe()
        assertNull(snapshot.gpu)
        assertNull(snapshot.processRssBytes)
    }
}
