import importlib.util
from pathlib import Path
import sys
import types

import pytest


ROOT = Path(__file__).resolve().parents[2]
EXPORTER_PATH = ROOT / "services" / "monitoring" / "scripts" / "nvidia_process_exporter.py"


def load_exporter(monkeypatch):
    """Load the exporter without requiring an NVIDIA driver on the test host."""

    pynvml = types.ModuleType("pynvml")

    class NVMLError(Exception):
        pass

    class NVMLErrorNotFound(NVMLError):
        pass

    class NVMLErrorNotSupported(NVMLError):
        pass

    pynvml.NVMLError = NVMLError
    pynvml.NVMLError_NotFound = NVMLErrorNotFound
    pynvml.NVMLError_NotSupported = NVMLErrorNotSupported
    for name in (
        "nvmlDeviceGetComputeRunningProcesses_v3",
        "nvmlDeviceGetCount",
        "nvmlDeviceGetGraphicsRunningProcesses_v3",
        "nvmlDeviceGetHandleByIndex",
        "nvmlDeviceGetName",
        "nvmlDeviceGetProcessUtilization",
        "nvmlDeviceGetUUID",
        "nvmlDeviceGetUtilizationRates",
        "nvmlInit",
    ):
        setattr(pynvml, name, lambda *args, **kwargs: None)

    monkeypatch.setitem(sys.modules, "pynvml", pynvml)
    spec = importlib.util.spec_from_file_location("nvidia_process_exporter_test", EXPORTER_PATH)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def test_process_sample_window_uses_nvml_microseconds(monkeypatch):
    exporter = load_exporter(monkeypatch)
    observed = []
    monkeypatch.setattr(exporter.time, "time", lambda: 1_700_000_000.0)
    monkeypatch.setattr(
        exporter,
        "nvmlDeviceGetProcessUtilization",
        lambda handle, since: observed.append(since) or [],
    )

    samples, supported = exporter.process_utilization_samples(object())

    assert samples == {}
    assert supported == 1
    assert observed == [1_700_000_000_000_000 - 30_000_000]


def test_namespace_attribution_scales_to_current_device_total(monkeypatch):
    exporter = load_exporter(monkeypatch)

    result = exporter.reconcile_namespace_utilization(
        {"game-stream": 20, "hermes": 10},
        device_util=3,
    )

    assert sum(result.values()) == pytest.approx(3)
    assert result["game-stream"] == pytest.approx(2)
    assert result["hermes"] == pytest.approx(1)


def test_namespace_attribution_assigns_unexplained_compute_to_host(monkeypatch):
    exporter = load_exporter(monkeypatch)

    result = exporter.reconcile_namespace_utilization(
        {"hermes": 1},
        device_util=3,
    )

    assert result == {"hermes": 1, "host": 2}


def test_zero_device_utilization_clears_stale_process_samples(monkeypatch):
    exporter = load_exporter(monkeypatch)

    result = exporter.reconcile_namespace_utilization(
        {"hermes": 40, "game-stream": 5},
        device_util=0,
    )

    assert result == {"hermes": 0, "game-stream": 0}