import psutil

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


class _DisplayStub:
    pass


def test_filter_processes_matches_name_case_insensitive() -> None:
    processes = [
        {"pid": 101, "name": "Python Train", "memory": 512.0},
        {"pid": 202, "name": "bash", "memory": 64.0},
    ]

    filtered = filter_processes(processes, "pYtHoN")

    assert [proc["pid"] for proc in filtered] == [101]


def test_filter_processes_matches_pid_substring() -> None:
    processes = [
        {"pid": 12345, "name": "worker", "memory": 512.0},
        {"pid": 98765, "name": "trainer", "memory": 256.0},
    ]

    filtered = filter_processes(processes, "234")

    assert [proc["pid"] for proc in filtered] == [12345]


def test_filter_processes_ignores_blank_query() -> None:
    processes = [
        {"pid": 1, "name": "init", "memory": 1.0},
        {"pid": 2, "name": "python", "memory": 2.0},
    ]

    filtered = filter_processes(processes, "   ")

    assert filtered == processes


def test_filter_processes_fielded_name_query() -> None:
    processes = [
        {"pid": 10, "name": "python train.py", "username": "alice"},
        {"pid": 20, "name": "tensorboard", "username": "alice"},
    ]

    filtered = filter_processes(processes, "name:python")

    assert [proc["pid"] for proc in filtered] == [10]


def test_filter_processes_fielded_pid_and_user_query() -> None:
    processes = [
        {"pid": 12345, "name": "worker", "username": "alice"},
        {"pid": 12399, "name": "worker", "username": "bob"},
        {"pid": 88888, "name": "worker", "username": "alice"},
    ]

    filtered = filter_processes(processes, "pid:123 user:ali")

    assert [proc["pid"] for proc in filtered] == [12345]


def test_filter_processes_fielded_cmd_and_plain_terms() -> None:
    processes = [
        {
            "pid": 200,
            "name": "python",
            "cmdline_summary": "python:train.py --dataset imagenet",
        },
        {
            "pid": 201,
            "name": "python",
            "cmdline_summary": "python:train.py --dataset cifar10",
        },
    ]

    filtered = filter_processes(processes, "cmd:train imagenet")

    assert [proc["pid"] for proc in filtered] == [200]


def test_filter_processes_gpu_boolean_query() -> None:
    processes = [
        {"pid": 300, "name": "python", "is_gpu_process": True},
        {"pid": 301, "name": "python", "is_gpu_process": False},
    ]

    only_gpu = filter_processes(processes, "gpu:true")
    only_non_gpu = filter_processes(processes, "gpu:false")

    assert [proc["pid"] for proc in only_gpu] == [300]
    assert [proc["pid"] for proc in only_non_gpu] == [301]


def test_filter_processes_unknown_field_falls_back_to_plain_token() -> None:
    processes = [
        {"pid": 400, "name": "owner:alice-worker"},
        {"pid": 401, "name": "owner:bob-worker"},
    ]

    filtered = filter_processes(processes, "owner:alice")

    assert [proc["pid"] for proc in filtered] == [400]


def test_filter_processes_all_terms_must_match() -> None:
    processes = [
        {"pid": 500, "name": "python train", "username": "alice"},
        {"pid": 501, "name": "python infer", "username": "bob"},
    ]

    filtered = filter_processes(processes, "python user:alice")

    assert [proc["pid"] for proc in filtered] == [500]


def test_gpu_panel_flattened_process_filter_respects_query() -> None:
    panel = GPUPanel(_DisplayStub())  # display is not used by this code path
    gpu_info = [
        (
            GPUInfo(
                index=0,
                name="GPU0",
                utilization=10.0,
                memory_used=1024.0,
                memory_total=8192.0,
                temperature=50.0,
                power_draw=100.0,
                power_limit=200.0,
                processes=[
                    {"pid": 111, "name": "python train.py", "memory": 2048.0},
                    {"pid": 222, "name": "tensorboard", "memory": 256.0},
                ],
            ),
            "nvidia",
        ),
        (
            GPUInfo(
                index=1,
                name="GPU1",
                utilization=15.0,
                memory_used=512.0,
                memory_total=8192.0,
                temperature=45.0,
                power_draw=80.0,
                power_limit=200.0,
                processes=[
                    {"pid": 111, "name": "python train.py", "memory": 3072.0},
                    {"pid": 333, "name": "bash", "memory": 32.0},
                ],
            ),
            "nvidia",
        ),
    ]

    filtered = panel.get_all_gpu_processes(gpu_info, "train")

    assert [proc["pid"] for proc in filtered] == [111]


def test_gpu_panel_selected_process_uses_filter_scope() -> None:
    panel = GPUPanel(_DisplayStub())
    gpu_info = [
        (
            GPUInfo(
                index=0,
                name="GPU0",
                utilization=20.0,
                memory_used=1024.0,
                memory_total=8192.0,
                temperature=55.0,
                power_draw=90.0,
                power_limit=200.0,
                processes=[
                    {"pid": 444, "name": "inference-service", "memory": 1024.0},
                    {"pid": 555, "name": "python eval.py", "memory": 512.0},
                ],
            ),
            "nvidia",
        )
    ]

    selected = panel.get_selected_process_by_index(gpu_info, 0, "eval")
    missing = panel.get_selected_process_by_index(gpu_info, 1, "eval")

    assert selected is not None
    assert selected["pid"] == 555
    assert missing is None


def test_gpu_panel_get_all_processes_assigns_global_npu_ids_across_vendors() -> None:
    panel = GPUPanel(_DisplayStub())
    npu_info = [
        (
            GPUInfo(
                index=0,
                name="Intel NPU",
                utilization=20.0,
                memory_used=0.0,
                memory_total=0.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[{"pid": 6601, "name": "rank-intel", "memory": 2.0}],
            ),
            "intel_npu",
        ),
        (
            GPUInfo(
                index=0,
                name="AMD NPU",
                utilization=25.0,
                memory_used=0.0,
                memory_total=0.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[{"pid": 6602, "name": "rank-amd", "memory": 1.0}],
            ),
            "amd_npu",
        ),
    ]

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

    assert by_pid[6601]["device"] == "NPU0"
    assert by_pid[6602]["device"] == "NPU1"
    assert by_pid[6601]["device_index"] == 0
    assert by_pid[6602]["device_index"] == 0
    assert by_pid[6601]["device_global_index"] == 0
    assert by_pid[6602]["device_global_index"] == 1


class _FakePsProcess:
    def __init__(self, pid: int) -> None:
        self._pid = pid

    def oneshot(self) -> "_FakePsProcess":
        return self

    def __enter__(self) -> "_FakePsProcess":
        return self

    def __exit__(self, exc_type: object, exc: object, tb: object) -> bool:
        del exc_type, exc, tb
        return False

    def ppid(self) -> int:
        return 42

    def cmdline(self) -> list[str]:
        if self._pid == 777:
            return ["python3", "/workspace/train.py", "--epochs", "1"]
        return ["python3", "/workspace/serve.py"]


def test_gpu_panel_enriches_runtime_name_from_cmdline(monkeypatch: object) -> None:
    panel = GPUPanel(_DisplayStub())
    gpu_info = [
        (
            GPUInfo(
                index=0,
                name="GPU0",
                utilization=20.0,
                memory_used=1024.0,
                memory_total=8192.0,
                temperature=55.0,
                power_draw=90.0,
                power_limit=200.0,
                processes=[{"pid": 777, "name": "python3", "memory": 512.0}],
            ),
            "nvidia",
        )
    ]

    monkeypatch.setattr(gpu_panel_module.psutil, "Process", _FakePsProcess)

    processes = panel.get_all_gpu_processes(gpu_info)

    assert len(processes) == 1
    assert processes[0]["name"] == "python:train.py"
    assert processes[0]["ppid"] == 42
    assert processes[0]["cmdline_summary"] == "python:train.py"
    assert "train.py" in processes[0]["cmdline"]


def test_gpu_panel_cmd_filter_uses_runtime_cmdline(monkeypatch: object) -> None:
    panel = GPUPanel(_DisplayStub())
    gpu_info = [
        (
            GPUInfo(
                index=0,
                name="GPU0",
                utilization=10.0,
                memory_used=1024.0,
                memory_total=8192.0,
                temperature=50.0,
                power_draw=100.0,
                power_limit=200.0,
                processes=[
                    {"pid": 777, "name": "python3", "memory": 512.0},
                    {"pid": 888, "name": "python3", "memory": 256.0},
                ],
            ),
            "nvidia",
        )
    ]

    monkeypatch.setattr(gpu_panel_module.psutil, "Process", _FakePsProcess)

    filtered = panel.get_all_gpu_processes(gpu_info, "cmd:train.py")

    assert [proc["pid"] for proc in filtered] == [777]


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

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


def test_gpu_panel_renders_fallback_util_for_npu_processes(monkeypatch: object) -> None:
    panel = GPUPanel(_RecordingDisplay())
    gpu = GPUInfo(
        index=0,
        name="AMD NPU",
        utilization=80.0,
        memory_used=0.0,
        memory_total=0.0,
        temperature=0.0,
        power_draw=0.0,
        power_limit=0.0,
        processes=[
            {"pid": 9101, "name": "rank0", "memory": 3.0},
            {"pid": 9102, "name": "rank1", "memory": 1.0},
        ],
    )

    class _NoRuntimeProc:
        def __init__(self, _pid: int) -> None:
            raise psutil.AccessDenied(pid=0)

    monkeypatch.setattr(gpu_panel_module.psutil, "Process", _NoRuntimeProc)
    monkeypatch.setattr(gpu_panel_module.curses, "color_pair", lambda _index: 0)
    monkeypatch.setattr(gpu_panel_module.curses, "A_REVERSE", 0)

    panel._render_processes(
        gpu=gpu,
        vendor="amd_npu",
        y=0,
        x=0,
        selected_pid=None,
        process_filter="",
    )

    lines = [entry[2] for entry in panel.display.calls]
    process_lines = [line for line in lines if "PID" in line and "|" in line]
    assert any("NPU0 60.0% |" in line for line in process_lines)
    assert any("NPU0 20.0% |" in line for line in process_lines)


def test_gpu_panel_header_renders_npu_driver_metadata(monkeypatch: object) -> None:
    panel = GPUPanel(_RecordingDisplay())
    gpu = GPUInfo(
        index=0,
        name="Intel AI Boost (0x7d1d)",
        utilization=0.0,
        memory_used=0.0,
        memory_total=0.0,
        temperature=0.0,
        power_draw=0.0,
        power_limit=0.0,
        processes=[],
        driver_version="1.2.3",
    )

    monkeypatch.setattr(gpu_panel_module.curses, "color_pair", lambda _index: 0)
    monkeypatch.setattr(gpu_panel_module.curses, "A_BOLD", 0)

    next_y = panel._render_gpu_header(
        gpu=gpu,
        y=5,
        x=2,
        vendor="intel_npu",
        device_label="NPU",
    )

    lines = [entry[2] for entry in panel.display.calls]
    assert any("NPU 0: Intel AI Boost (0x7d1d)" in line for line in lines)
    assert any("Driver: 1.2.3" in line for line in lines)
    assert next_y == 7
