import aitop.ui.components.gpu_panel as gpu_panel_module
from aitop.core.gpu.base import GPUInfo
from aitop.ui.components.gpu_panel import GPUPanel


class _RecordingDisplay:
    def __init__(self) -> None:
        self.calls: list[tuple[int, int, str, int]] = []
        self.width = 140

    def safe_addstr(self, y: int, x: int, text: str, attr: int = 0) -> None:
        self.calls.append((y, x, text, attr))

    def create_bar(self, _value: float, width: int) -> tuple[str, int]:
        return ("#" * min(max(width, 0), 5), 0)

    def get_color(self, _value: float) -> int:
        return 0


def test_gpu_panel_memory_bar_handles_missing_total() -> None:
    display = _RecordingDisplay()
    panel = GPUPanel(display)
    gpu = GPUInfo(
        index=0,
        name="NVIDIA A100",
        utilization=12.0,
        memory_used=0.0,
        memory_total=0.0,
        temperature=0.0,
        power_draw=0.0,
        power_limit=0.0,
        processes=[],
    )

    panel._render_memory_bar(gpu, y=0, x=0, width=20)

    assert len(display.calls) == 1
    assert "0.0GB" in display.calls[0][2]


def test_gpu_panel_memory_bar_marks_shared_memory_model() -> None:
    display = _RecordingDisplay()
    panel = GPUPanel(display)
    gpu = GPUInfo(
        index=0,
        name="Intel iGPU",
        utilization=12.0,
        memory_used=1024.0,
        memory_total=8192.0,
        temperature=0.0,
        power_draw=0.0,
        power_limit=0.0,
        processes=[],
        memory_model="shared",
    )

    panel._render_memory_bar(gpu, y=0, x=0, width=20)

    assert len(display.calls) == 1
    assert "Mem~" in display.calls[0][2]


def test_gpu_panel_skips_temperature_and_power_when_missing() -> None:
    display = _RecordingDisplay()
    panel = GPUPanel(display)
    gpu = GPUInfo(
        index=0,
        name="NVIDIA A100",
        utilization=12.0,
        memory_used=512.0,
        memory_total=0.0,
        temperature=0.0,
        power_draw=200.0,
        power_limit=0.0,
        processes=[],
    )

    y_after_temp = panel._render_temperature(gpu, y=0, x=0)
    y_after_power = panel._render_power(gpu, y=y_after_temp, x=0)

    assert y_after_temp == 0
    assert y_after_power == 0
    assert display.calls == []


def test_gpu_panel_process_tree_converts_bytes_memory() -> None:
    display = _RecordingDisplay()
    panel = GPUPanel(display)
    proc = {
        "pid": 1234,
        "name": "trainer",
        "memory": 2 * 1024 * 1024 * 1024,
        "device": "GPU0",
    }

    gpu_panel_module.curses.color_pair = lambda _index: 0  # type: ignore[assignment]
    gpu_panel_module.curses.A_REVERSE = 0  # type: ignore[assignment]

    panel._render_process_tree_node(
        proc=proc,
        y=0,
        x=0,
        selected_pid=None,
        children_map={},
        is_last=True,
        prefix="",
        device_label="GPU",
    )

    assert any("2.0GB" in line[2] for line in display.calls)


def test_gpu_panel_build_context_processes_handles_missing_pid() -> None:
    display = _RecordingDisplay()
    panel = GPUPanel(display)
    gpu = GPUInfo(
        index=0,
        name="NVIDIA A100",
        utilization=0.0,
        memory_used=0.0,
        memory_total=0.0,
        temperature=0.0,
        power_draw=0.0,
        power_limit=0.0,
        processes=[{"pid": "not-int", "name": "unknown", "memory": 0.0}],
    )

    processes = panel._build_context_processes(gpu, "nvidia", runtime_cache={})

    assert processes[0]["name"] == "unknown"
    assert processes[0]["device"] == "GPU0"


def test_gpu_panel_process_tree_handles_missing_memory() -> None:
    display = _RecordingDisplay()
    panel = GPUPanel(display)
    proc = {
        "pid": 5555,
        "name": "worker",
        "device": "GPU0",
    }

    gpu_panel_module.curses.color_pair = lambda _index: 0  # type: ignore[assignment]
    gpu_panel_module.curses.A_REVERSE = 0  # type: ignore[assignment]

    panel._render_process_tree_node(
        proc=proc,
        y=0,
        x=0,
        selected_pid=None,
        children_map={},
        is_last=True,
        prefix="",
        device_label="GPU",
    )

    assert any("0.0MB" in line[2] for line in display.calls)


def test_gpu_panel_process_tree_clamps_utilization_percent() -> None:
    display = _RecordingDisplay()
    panel = GPUPanel(display)
    proc = {
        "pid": 9999,
        "name": "worker",
        "memory": 1024.0,
        "device": "GPU0",
        "sm_util": 2500.0,
    }

    gpu_panel_module.curses.color_pair = lambda _index: 0  # type: ignore[assignment]
    gpu_panel_module.curses.A_REVERSE = 0  # type: ignore[assignment]

    panel._render_process_tree_node(
        proc=proc,
        y=0,
        x=0,
        selected_pid=None,
        children_map={},
        is_last=True,
        prefix="",
        device_label="GPU",
    )

    assert any("GPU0 100.0%" in line[2] for line in display.calls)


def test_get_all_gpu_processes_assigns_global_gpu_ids_across_vendors() -> None:
    panel = GPUPanel(_RecordingDisplay())
    gpu_info = [
        (
            GPUInfo(
                index=0,
                name="NVIDIA H100",
                utilization=70.0,
                memory_used=12000.0,
                memory_total=80000.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[{"pid": 1111, "name": "train-a", "memory": 1024.0}],
            ),
            "nvidia",
        ),
        (
            GPUInfo(
                index=0,
                name="AMD MI300",
                utilization=60.0,
                memory_used=8000.0,
                memory_total=64000.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[{"pid": 2222, "name": "train-b", "memory": 2048.0}],
            ),
            "amd",
        ),
    ]

    processes = panel.get_all_gpu_processes(gpu_info)
    by_pid = {proc["pid"]: proc for proc in processes}

    assert by_pid[1111]["device"] == "GPU0"
    assert by_pid[2222]["device"] == "GPU1"
    assert by_pid[1111]["device_index"] == 0
    assert by_pid[2222]["device_index"] == 0
    assert by_pid[1111]["device_global_index"] == 0
    assert by_pid[2222]["device_global_index"] == 1
