"""Pure helpers in vram_arbitrator: NVML throttle-bit decoding and PID attribution.

Deliberately excludes instant_free_ollama_vram, the AutoArbitrator yield/purge paths and
the SSE broker — that contract is in flux.
"""
import vram_arbitrator as va


def test_decode_throttle_reasons_empty_when_no_bits_set():
    assert va.decode_throttle_reasons(0) == []


def test_decode_throttle_reasons_maps_each_known_bit():
    """Every mask in the table must decode to exactly its own name in isolation."""
    for mask, name in va.THROTTLE_REASONS.items():
        assert va.decode_throttle_reasons(mask) == [name]


def test_decode_throttle_reasons_decodes_combined_bits():
    """Real NVML samples set several bits at once; all of them must come back."""
    bits = 0x20 | 0x40  # sw_thermal_slowdown | hw_thermal_slowdown
    assert set(va.decode_throttle_reasons(bits)) == {"sw_thermal_slowdown", "hw_thermal_slowdown"}


def test_decode_throttle_reasons_ignores_unknown_bits():
    """An undocumented bit from a future driver must not crash or invent a reason."""
    assert va.decode_throttle_reasons(0x8000_0000) == []


def test_hard_throttle_names_match_thermal_governor_expectations():
    """thermal_governor escalates on a fixed set of reason strings produced here.
    If a name is renamed in one module and not the other the governor silently stops
    reacting to hardware slowdowns, so pin the shared vocabulary."""
    import thermal_governor as tg
    assert tg.HARD_THROTTLES <= set(va.THROTTLE_REASONS.values())


class _FakeProc:
    def __init__(self, name, cmdline):
        self._name = name
        self._cmdline = cmdline

    def name(self):
        return self._name

    def cmdline(self):
        return self._cmdline


def _patch_proc(monkeypatch, proc):
    monkeypatch.setattr(va.psutil, "Process", lambda pid: proc)


def test_classify_pid_detects_ollama_by_process_name(monkeypatch):
    _patch_proc(monkeypatch, _FakeProc("ollama", ["/usr/local/bin/ollama", "serve"]))
    assert va._classify_pid(1234) == "ollama"


def test_classify_pid_detects_ollama_runner_by_cmdline(monkeypatch):
    """Ollama's model runner is a separate llama-server process; its VRAM is Ollama's."""
    _patch_proc(monkeypatch, _FakeProc("llama-server",
                                       ["/usr/lib/ollama/llama-server", "--model", "blob"]))
    assert va._classify_pid(1234) == "ollama"


def test_classify_pid_detects_comfyui(monkeypatch):
    _patch_proc(monkeypatch, _FakeProc("python3", ["python3", "/opt/ComfyUI/main.py", "--listen"]))
    assert va._classify_pid(1234) == "comfy"


def test_classify_pid_unknown_process_is_other(monkeypatch):
    _patch_proc(monkeypatch, _FakeProc("Xorg", ["/usr/lib/xorg/Xorg", ":8"]))
    assert va._classify_pid(1234) == "other"


def test_classify_pid_returns_other_when_process_vanished(monkeypatch):
    """PIDs are read from NVML and can exit before psutil looks them up; that is normal
    and must not raise inside the 20 ms VRAM poll loop."""
    def _boom(pid):
        raise va.psutil.NoSuchProcess(pid)
    monkeypatch.setattr(va.psutil, "Process", _boom)
    assert va._classify_pid(999999) == "other"