import json
import subprocess
import time
from pathlib import Path

from aitop.core.gpu.amd import AMDGPUMonitor
from aitop.core.gpu.intel import IntelGPUMonitor
from aitop.core.gpu.nvidia import NvidiaGPUMonitor

FIXTURE_DIR = Path(__file__).parent / "fixtures" / "gpu"


def _read_fixture(name: str) -> str:
    return (FIXTURE_DIR / name).read_text(encoding="utf-8")


def test_nvidia_parses_basic_and_advanced_fixture_rows() -> None:
    monitor = NvidiaGPUMonitor()

    monitor._parse_gpu_data("basic", _read_fixture("nvidia_basic.csv"))
    monitor._parse_gpu_data("advanced", _read_fixture("nvidia_advanced.csv"))

    assert set(monitor._data_cache["basic"].keys()) == {0, 1}
    assert set(monitor._data_cache["advanced"].keys()) == {0, 1}

    gpu0 = monitor._data_cache["basic"][0]
    gpu1 = monitor._data_cache["basic"][1]
    advanced0 = monitor._data_cache["advanced"][0]
    advanced1 = monitor._data_cache["advanced"][1]

    assert gpu0["name"] == "NVIDIA A100-SXM4-40GB"
    assert gpu0["utilization"] == 98.0
    assert gpu0["memory_used"] == 38210.0
    assert gpu0["memory_total"] == 40960.0
    assert gpu1["utilization"] == 0.0
    assert advanced0["temperature"] == 72.0
    assert advanced0["power_draw"] == 285.5
    assert advanced1["power_limit"] == 0.0


def test_nvidia_parses_process_and_pmon_fixtures() -> None:
    monitor = NvidiaGPUMonitor()

    monitor._parse_gpu_data("basic", _read_fixture("nvidia_basic.csv"))
    monitor._parse_process_data(_read_fixture("nvidia_processes.csv"))
    monitor._parse_process_util_data(_read_fixture("nvidia_pmon.txt"))

    process_entries = monitor._data_cache["processes"]
    util_entries = monitor._data_cache["process_util"]

    assert len(process_entries) == 2
    assert {entry["pid"] for entry in process_entries} == {910001, 910002}
    assert process_entries[0]["gpu_index"] == 0
    assert process_entries[1]["gpu_index"] == 1
    assert util_entries[910001]["sm_util"] == 95.0
    assert util_entries[910001]["enc_util"] == 0.0
    assert util_entries[910002]["sm_util"] == 0.0
    assert util_entries[910002]["mem_util"] == 4.0
    assert monitor._data_cache["process_util_by_gpu"][0] == 95.0
    assert monitor._data_cache["process_util_by_gpu"][1] == 0.0


def test_nvidia_parses_pmon_with_mig_columns() -> None:
    monitor = NvidiaGPUMonitor()
    pmon_output = (
        "# gpu gi ci pid type sm mem enc dec command\n"
        "# Idx Gi Ci  # C/G  %   %   %   %  name\n"
        "0 3 2 910010 C 70 10 0 0 python\n"
        "0 3 2 910011 C 40  8 - - worker\n"
        "1 - - 910012 C -   4 0 0 tritonserver\n"
    )

    monitor._parse_process_util_data(pmon_output)

    assert monitor._data_cache["process_util"][910010]["sm_util"] == 70.0
    assert monitor._data_cache["process_util"][910011]["sm_util"] == 40.0
    assert monitor._data_cache["process_util"][910012]["sm_util"] == 0.0
    assert monitor._data_cache["process_util_by_gpu"][0] == 100.0
    assert monitor._data_cache["process_util_by_gpu"][1] == 0.0


def test_nvidia_get_gpu_info_builds_complete_objects_from_fixture_data(
    monkeypatch: object,
) -> None:
    monitor = NvidiaGPUMonitor()
    monitor._cached_smi_path_initialized = True
    monitor._cached_smi_path = Path("nvidia-smi")

    monitor._parse_gpu_data("basic", _read_fixture("nvidia_basic.csv"))
    monitor._parse_gpu_data("advanced", _read_fixture("nvidia_advanced.csv"))
    monitor._parse_process_data(_read_fixture("nvidia_processes.csv"))
    monitor._parse_process_util_data(_read_fixture("nvidia_pmon.txt"))

    monkeypatch.setattr(monitor, "_update_cache_if_needed", lambda: False)
    monkeypatch.setattr(monitor, "_get_driver_version", lambda: "550.54.14")
    monkeypatch.setattr(monitor, "_get_cuda_version", lambda: "12.4")

    gpus = monitor.get_gpu_info()

    assert [gpu.index for gpu in gpus] == [0, 1]
    assert gpus[0].driver_version == "550.54.14"
    assert gpus[0].cuda_version == "12.4"
    assert gpus[0].processes[0]["pid"] == 910001
    assert gpus[0].processes[0]["sm_util"] == 95.0
    assert gpus[1].processes[0]["pid"] == 910002
    assert gpus[1].processes[0]["sm_util"] == 0.0


def test_nvidia_get_gpu_info_falls_back_to_pmon_gpu_utilization(
    monkeypatch: object,
) -> None:
    monitor = NvidiaGPUMonitor()
    monitor._cached_smi_path_initialized = True
    monitor._cached_smi_path = Path("nvidia-smi")

    monitor._parse_gpu_data("basic", _read_fixture("nvidia_basic.csv"))
    monitor._parse_process_util_data(
        "# gpu pid type sm mem enc dec command\n"
        "# Idx # C/G  %  %   %   %  name\n"
        "1 920001 C 37 12 0 0 python\n"
    )

    monkeypatch.setattr(monitor, "_update_cache_if_needed", lambda: False)
    monkeypatch.setattr(monitor, "_get_driver_version", lambda: "")
    monkeypatch.setattr(monitor, "_get_cuda_version", lambda: "")

    gpus = monitor.get_gpu_info()
    util_by_index = {gpu.index: gpu.utilization for gpu in gpus}

    assert util_by_index[0] == 98.0
    assert util_by_index[1] == 37.0


def test_nvidia_quick_metrics_falls_back_to_cached_pmon_utilization() -> None:
    monitor = NvidiaGPUMonitor()
    monitor._smi_path_initialized = True
    monitor._smi_path = Path("nvidia-smi")

    monitor._parse_gpu_data("basic", _read_fixture("nvidia_basic.csv"))
    monitor._parse_process_util_data(
        "# gpu pid type sm mem enc dec command\n"
        "# Idx # C/G  %  %   %   %  name\n"
        "1 930001 C 42 11 0 0 torchrun\n"
    )
    monitor._last_updates["basic"] = time.time()

    metrics = monitor.get_quick_metrics()

    assert metrics[0]["utilization"] == 98.0
    assert metrics[1]["utilization"] == 42.0


def test_nvidia_update_cache_if_needed_populates_all_query_types(
    monkeypatch: object,
) -> None:
    monitor = NvidiaGPUMonitor()
    monitor._smi_path_initialized = True
    monitor._smi_path = Path("nvidia-smi")

    def fake_batch_runner(commands: list[tuple[str, list[str]]]) -> dict[str, str]:
        query_type = commands[0][0]
        payloads = {
            "basic": _read_fixture("nvidia_basic.csv"),
            "advanced": _read_fixture("nvidia_advanced.csv"),
            "processes": _read_fixture("nvidia_processes.csv"),
            "process_util": _read_fixture("nvidia_pmon.txt"),
        }
        return {query_type: payloads.get(query_type, "")}

    monkeypatch.setattr(monitor, "_run_smi_batch_command", fake_batch_runner)

    updated = monitor._update_cache_if_needed()

    assert updated is True
    assert monitor._data_cache["basic"][0]["utilization"] == 98.0
    assert monitor._data_cache["advanced"][0]["temperature"] == 72.0
    assert len(monitor._data_cache["processes"]) == 2
    assert monitor._data_cache["process_util"][910001]["sm_util"] == 95.0


def test_nvidia_parses_advanced_row_without_clock_column() -> None:
    monitor = NvidiaGPUMonitor()
    monitor._parse_gpu_data("advanced", "0, 68, 210.5, 300.0")

    assert monitor._data_cache["advanced"][0]["temperature"] == 68.0
    assert monitor._data_cache["advanced"][0]["power_draw"] == 210.5
    assert monitor._data_cache["advanced"][0]["power_limit"] == 300.0
    assert monitor._data_cache["advanced"][0]["clock"] == 0.0


def test_nvidia_parses_process_row_without_name_field() -> None:
    monitor = NvidiaGPUMonitor()
    monitor._gpu_uuids[0] = "GPU-aaaaaaaa-bbbb-cccc-dddd-eeeeeeee0000"

    monitor._parse_process_data(
        "GPU-aaaaaaaa-bbbb-cccc-dddd-eeeeeeee0000, 99999999, 2048"
    )

    assert len(monitor._data_cache["processes"]) == 1
    assert monitor._data_cache["processes"][0]["pid"] == 99999999
    assert monitor._data_cache["processes"][0]["gpu_index"] == 0
    assert monitor._data_cache["processes"][0]["memory"] == 2048.0
    assert monitor._data_cache["processes"][0]["name"] == "unknown"


def test_nvidia_process_parser_skips_invalid_pid_and_defaults_bad_memory() -> None:
    monitor = NvidiaGPUMonitor()
    monitor._gpu_uuids[0] = "GPU-uuid-0"

    monitor._parse_process_data(
        "\n".join(
            [
                "GPU-uuid-0, not-a-pid, 1024, python",
                "GPU-uuid-0, 4242, not-a-number, trainer",
            ]
        )
    )

    assert len(monitor._data_cache["processes"]) == 1
    assert monitor._data_cache["processes"][0]["pid"] == 4242
    assert monitor._data_cache["processes"][0]["memory"] == 0.0


def test_nvidia_pmon_handles_na_values() -> None:
    monitor = NvidiaGPUMonitor()
    pmon_output = (
        "# gpu pid type sm mem enc dec command\n"
        "# Idx # C/G % % % % name\n"
        "0 555 C N/A - - - python\n"
    )

    monitor._parse_process_util_data(pmon_output)

    assert monitor._data_cache["process_util"][555]["sm_util"] == 0.0
    assert monitor._data_cache["process_util"][555]["mem_util"] == 0.0
    assert monitor._data_cache["process_util_by_gpu"][0] == 0.0


def test_nvidia_pmon_clamps_out_of_range_values() -> None:
    monitor = NvidiaGPUMonitor()
    pmon_output = (
        "# gpu pid type sm mem enc dec command\n"
        "# Idx # C/G % % % % name\n"
        "0 777 C 1500 600 - - python\n"
    )

    monitor._parse_process_util_data(pmon_output)

    assert monitor._data_cache["process_util"][777]["sm_util"] == 100.0
    assert monitor._data_cache["process_util"][777]["mem_util"] == 100.0
    assert monitor._data_cache["process_util_by_gpu"][0] == 100.0


def test_intel_parses_json_stream_with_truncated_tail() -> None:
    monitor = IntelGPUMonitor()
    output = (
        '{"card": 0, "engines": {"Render/3D/0": {"busy": 37.5}}},'
        '{"card": 0, "engines": {"Render/3D/0": {"busy":'
    )

    samples = monitor._parse_json_samples(output)

    assert len(samples) == 1
    assert samples[0]["card"] == 0
    assert monitor._extract_engine_utilization(samples[0]) == 37.5


def test_intel_extract_processes_keeps_highest_memory_per_pid() -> None:
    monitor = IntelGPUMonitor()
    sample = {
        "clients": [
            {"pid": "1001", "name": "job", "memory": 64},
            {"pid": 1001, "process": "job", "memory_used": 128},
            {"pid": "bad", "name": "skip", "memory": 256},
        ]
    }

    processes = monitor._extract_processes(sample)

    assert processes == [{"pid": 1001, "name": "job", "memory": 128.0}]


def test_intel_parses_gpu_info_from_json_sample_and_device_list() -> None:
    monitor = IntelGPUMonitor()
    devices = [{"index": 1, "name": "Intel Arc A770"}]
    output = (
        '{"card":"card1","engines":{"Render/3D/0":{"busy":42.0}},'
        '"clients":[{"pid":1234,"name":"python",'
        '"memory":{"value":1073741824,"unit":"B"}}]}'
    )

    gpu_info = monitor._parse_gpu_info(output, devices=devices)

    assert len(gpu_info) == 1
    assert gpu_info[0]["index"] == 1
    assert gpu_info[0]["name"] == "Intel Arc A770"
    assert gpu_info[0]["utilization"] == 42.0
    assert gpu_info[0]["processes"][0]["pid"] == 1234
    assert gpu_info[0]["processes"][0]["memory"] == 1024.0


def test_intel_device_list_dedupes_card_and_render_nodes() -> None:
    monitor = IntelGPUMonitor()
    output = (
        "/dev/dri/card0: Intel Graphics\n"
        "/dev/dri/renderD128: Intel Graphics Render Node\n"
        "List of devices\n"
    )

    devices = monitor._parse_device_list(output)

    assert devices == [{"index": 0, "name": "Intel Graphics"}]


def test_intel_build_gpu_records_collapses_unindexed_samples_to_single_device() -> None:
    monitor = IntelGPUMonitor()
    samples = [
        {"engines": {"Render/3D/0": {"busy": 11.0}}},
        {"engines": {"Render/3D/0": {"busy": 27.5}}},
    ]

    records = monitor._build_gpu_records(samples=samples, devices=[])

    assert len(records) == 1
    assert records[0]["index"] == 0
    assert records[0]["utilization"] == 27.5


def test_intel_run_smi_command_returns_partial_output_on_timeout(
    monkeypatch: object,
) -> None:
    monitor = IntelGPUMonitor()

    def _timeout(*args: object, **kwargs: object) -> object:
        del kwargs
        raise subprocess.TimeoutExpired(
            cmd=args[0],
            timeout=5,
            output='{"card":0,"engines":{"Render/3D/0":{"busy":12.0}}}',
        )

    monkeypatch.setattr(subprocess, "run", _timeout)

    result = monitor._run_smi_command(["intel_gpu_top", "-J", "-s", "1000"])

    assert result == '{"card":0,"engines":{"Render/3D/0":{"busy":12.0}}}'


def test_intel_resolve_memory_metrics_prefers_sysfs_shared_totals(
    monkeypatch: object,
) -> None:
    monitor = IntelGPUMonitor()

    def _fake_sysfs(path: Path) -> int:
        if path.name == "mem_info_system_total":
            return 8 * 1024 * 1024 * 1024
        if path.name == "mem_info_system_used":
            return 2 * 1024 * 1024 * 1024
        return 0

    monkeypatch.setattr(monitor, "_read_sysfs_int", _fake_sysfs)

    used_mb, total_mb, model, source, confidence = monitor._resolve_memory_metrics(
        index=0,
        memory_used_mb=0.0,
        memory_total_mb=0.0,
    )

    assert used_mb == 2048.0
    assert total_mb == 8192.0
    assert model == "shared"
    assert source == "sysfs:mem_info_system"
    assert confidence == "direct"


def test_amd_parses_showpids_verbose_fixture() -> None:
    monitor = AMDGPUMonitor()

    processes = monitor._parse_processes(_read_fixture("amd_showpids_verbose.txt"))

    assert len(processes) == 2
    assert processes[0]["pid"] == 920001
    assert processes[0]["gpu_ids"] == [0]
    assert processes[0]["vram_used"] == 4096.0
    assert processes[0]["cu_occupancy"] == 72.0
    assert processes[1]["gpu_ids"] == [0, 1]
    assert processes[1]["cu_occupancy"] is None


def test_amd_shared_gtt_memory_info_reads_sysfs(monkeypatch: object) -> None:
    monitor = AMDGPUMonitor()

    def _fake_sysfs(path: Path) -> int:
        if path.name == "mem_info_gtt_total":
            return 4 * 1024 * 1024 * 1024
        if path.name == "mem_info_gtt_used":
            return 1024 * 1024 * 1024
        return 0

    monkeypatch.setattr(monitor, "_read_sysfs_int", _fake_sysfs)

    used_mb, total_mb = monitor._get_shared_gtt_memory_info(0)

    assert used_mb == 1024.0
    assert total_mb == 4096.0


def test_amd_parses_batch_json_metrics_fixture() -> None:
    monitor = AMDGPUMonitor()
    batch_data = json.loads(_read_fixture("amd_batch_metrics.json"))

    used_mb, total_mb = monitor._parse_json_memory(batch_data, gpu_index=0)
    temperature = monitor._parse_json_temperature(batch_data, gpu_index=0)
    power = monitor._parse_json_power(batch_data, gpu_index=0)
    utilization = monitor._parse_json_utilization(batch_data, gpu_index=0)

    assert used_mb == 8192.0
    assert total_mb == 16384.0
    assert temperature == 71.5
    assert power == 265.4
    assert utilization == 88.0

    missing_used, missing_total = monitor._parse_json_memory(batch_data, gpu_index=99)
    assert missing_used == 0.0
    assert missing_total == 0.0


def test_amd_parses_amd_smi_style_metric_payload() -> None:
    monitor = AMDGPUMonitor()
    batch_data = {
        "gpu0": {
            "average_gfx_activity": 82,
            "vram_total_memory_b": 17179869184,
            "vram_total_used_memory_b": 4294967296,
            "junction_temperature_c": 68.5,
            "average_socket_power_w": 214.2,
        }
    }

    used_mb, total_mb = monitor._parse_json_memory(batch_data, gpu_index=0)
    temperature = monitor._parse_json_temperature(batch_data, gpu_index=0)
    power = monitor._parse_json_power(batch_data, gpu_index=0)
    utilization = monitor._parse_json_utilization(batch_data, gpu_index=0)

    assert used_mb == 4096.0
    assert total_mb == 16384.0
    assert temperature == 68.5
    assert power == 214.2
    assert utilization == 82.0


def test_amd_parses_amd_smi_process_payload() -> None:
    monitor = AMDGPUMonitor()
    process_data = {
        "gpu0": {
            "processes": [
                {
                    "pid": 930001,
                    "process_name": "python",
                    "gpu_index": 0,
                    "vram_used_mb": 1536,
                    "cu_occupancy": 57.0,
                },
                {
                    "pid": 930002,
                    "process_name": "torchrun",
                    "gpu_index": "gpu0",
                    "vram_used_b": 1073741824,
                },
            ]
        }
    }

    processes = monitor._parse_amd_smi_processes(process_data, gpu_index=0)
    assert len(processes) == 2

    assert processes[0]["pid"] == 930001
    assert processes[0]["gpu_ids"] == [0]
    assert processes[0]["vram_used"] == 1536.0
    assert processes[0]["cu_occupancy"] == 57.0

    assert processes[1]["pid"] == 930002
    assert processes[1]["gpu_ids"] == [0]
    assert processes[1]["vram_used"] == 1024.0


def test_amd_smi_processes_parse_gpu_ids_and_memory_units() -> None:
    monitor = AMDGPUMonitor()
    payload = [
        {
            "pid": "4321",
            "process name": "trainer",
            "gpu": "gpu0, gpu1",
            "vram used b": 1024 * 1024 * 2048,
            "gpu usage": "75%",
        },
        {
            "pid": 4322,
            "cmdline": "worker",
            "gpu": "card2",
            "vram used": "512 MB",
        },
        {"pid": "not-a-number", "name": "skip"},
    ]

    processes = monitor._parse_amd_smi_processes(payload)

    assert len(processes) == 2
    assert processes[0]["pid"] == 4321
    assert processes[0]["gpu_ids"] == [0, 1]
    assert round(processes[0]["memory"], 1) == 2048.0
    assert processes[0]["cu_occupancy"] == 75.0
    assert processes[1]["pid"] == 4322
    assert processes[1]["gpu_ids"] == [2]
    assert processes[1]["memory"] == 512.0


def test_amd_parses_nested_value_unit_metric_payload() -> None:
    monitor = AMDGPUMonitor()
    batch_data = [
        {
            "gpu": 0,
            "usage": {
                "average_gfx_activity": {"value": 91, "unit": "%"},
            },
            "vram_usage": {
                "vram_total": {"value": 17179869184, "unit": "B"},
                "vram_used": {"value": 4294967296, "unit": "B"},
            },
            "temperature": {"junction": {"value": 72.0, "unit": "C"}},
            "power": {"socket_power": {"value": 220, "unit": "W"}},
        }
    ]

    used_mb, total_mb = monitor._parse_json_memory(batch_data, gpu_index=0)
    temperature = monitor._parse_json_temperature(batch_data, gpu_index=0)
    power = monitor._parse_json_power(batch_data, gpu_index=0)
    utilization = monitor._parse_json_utilization(batch_data, gpu_index=0)

    assert used_mb == 4096.0
    assert total_mb == 16384.0
    assert temperature == 72.0
    assert power == 220.0
    assert utilization == 91.0


def test_amd_parses_nested_process_mem_usage_payload() -> None:
    monitor = AMDGPUMonitor()
    process_data = [
        {
            "gpu": 0,
            "process_list": [
                {
                    "process_info": {
                        "pid": 930100,
                        "name": "python",
                        "mem_usage": {"value": 1073741824, "unit": "B"},
                        "gpu_usage": {"value": 45.0, "unit": "%"},
                    }
                }
            ],
        }
    ]

    processes = monitor._parse_amd_smi_processes(process_data, gpu_index=0)

    assert len(processes) == 1
    assert processes[0]["pid"] == 930100
    assert processes[0]["gpu_ids"] == [0]
    assert processes[0]["vram_used"] == 1024.0
    assert processes[0]["cu_occupancy"] == 45.0
