import pytest

import aitop.core.process.monitor as monitor_module
from aitop.core.process.monitor import AIProcessMonitor


class _FakeProcess:
    def __init__(self, pid: int, name: str) -> None:
        self.pid = pid
        self.info = {"pid": pid, "name": name}


@pytest.mark.parametrize(
    ("proc_name", "cmdline"),
    [
        ("vllm", "vllm serve meta-llama/Llama-3.1-8B-Instruct"),
        (
            "python",
            "python -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B",
        ),
        ("text-generation-launcher", "text-generation-launcher --model-id mistral"),
        ("sglang", "python -m sglang.launch_server --model-path Qwen/Qwen2.5-7B"),
        ("trtllm-serve", "trtllm-serve --model /models/llama"),
        ("tritonserver", "tritonserver --model-repository=/models"),
        ("llama-server", "llama-server -m /models/model.gguf"),
        ("ollama", "ollama serve"),
        ("deepspeed", "deepspeed --num_gpus=2 train.py"),
        ("accelerate", "accelerate launch train.py"),
        ("torchrun", "torchrun --nproc-per-node=4 train.py"),
        ("python", "python:finetune_llm.py"),
    ],
)
def test_is_ai_process_matches_known_runtime_signatures(
    proc_name: str, cmdline: str
) -> None:
    monitor = AIProcessMonitor()

    assert monitor.is_ai_process(proc_name, cmdline) is True


@pytest.mark.parametrize(
    ("proc_name", "cmdline"),
    [
        ("chrome", "chrome --enable-gpu-rasterization --use-gl=desktop"),
        ("nvidia-smi", "nvidia-smi"),
        ("slurmstepd", "slurmstepd: [1234]"),
        ("systemd", "/usr/lib/systemd/systemd --user"),
        ("python", "python app.py --mode serve --port 8080"),
        ("python", "python:backup_metrics.py"),
    ],
)
def test_is_ai_process_rejects_common_false_positives(
    proc_name: str, cmdline: str
) -> None:
    monitor = AIProcessMonitor()

    assert monitor.is_ai_process(proc_name, cmdline) is False


def test_is_ai_process_fallback_patterns_are_strict(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    def _raise(_: object) -> list[str]:
        raise RuntimeError("boom")

    monkeypatch.setattr(monitor_module, "load_process_patterns", _raise)

    monitor = monitor_module.AIProcessMonitor()

    assert monitor.is_ai_process("vllm", "vllm serve meta-llama/Llama-3.1-8B") is True
    assert (
        monitor.is_ai_process(
            "chrome", "chrome --enable-gpu-rasterization --use-gl=desktop"
        )
        is False
    )


def test_is_ai_process_cache_is_cmdline_aware_for_python() -> None:
    monitor = AIProcessMonitor()

    assert (
        monitor.is_ai_process(
            "python", "python -m vllm.entrypoints.openai.api_server --model qwen"
        )
        is True
    )
    assert (
        monitor.is_ai_process("python", "python app.py --mode serve --port 8080")
        is False
    )


def test_get_ai_processes_excludes_unmatched_gpu_processes(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monitor = AIProcessMonitor()
    fake_proc = _FakeProcess(pid=1234, name="chrome")

    monitor._last_full_scan = monitor_module.time.time()

    monkeypatch.setattr(
        monitor_module.psutil, "process_iter", lambda _attrs: [fake_proc]
    )
    monkeypatch.setattr(monitor_module.psutil, "pids", lambda: [1234])
    monkeypatch.setattr(
        monitor,
        "_get_process_cmdline",
        lambda _proc: "chrome --enable-gpu-rasterization",
    )
    monkeypatch.setattr(monitor, "is_ai_process", lambda _name, _cmdline: False)
    monkeypatch.setattr(
        monitor,
        "_get_process_info",
        lambda _proc, _now, include_extended: {
            "pid": 1234,
            "name": "chrome",
            "cmdline": "chrome --enable-gpu-rasterization",
            "extended": include_extended,
        },
    )

    ai_processes = monitor.get_ai_processes(
        gpu_processes=[{"pid": 1234}], force_refresh=False
    )

    assert ai_processes == []


def test_get_ai_processes_keeps_gpu_flag_for_matched_ai_processes(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monitor = AIProcessMonitor()
    fake_proc = _FakeProcess(pid=4321, name="vllm")

    monitor._last_full_scan = monitor_module.time.time()

    monkeypatch.setattr(
        monitor_module.psutil, "process_iter", lambda _attrs: [fake_proc]
    )
    monkeypatch.setattr(monitor_module.psutil, "pids", lambda: [4321])
    monkeypatch.setattr(
        monitor, "_get_process_cmdline", lambda _proc: "vllm serve qwen"
    )
    monkeypatch.setattr(monitor, "is_ai_process", lambda _name, _cmdline: True)
    monkeypatch.setattr(
        monitor,
        "_get_process_info",
        lambda _proc, _now, include_extended: {
            "pid": 4321,
            "name": "vllm",
            "cmdline": "vllm serve qwen",
            "extended": include_extended,
        },
    )

    ai_processes = monitor.get_ai_processes(
        gpu_processes=[{"pid": 4321}], force_refresh=False
    )

    assert len(ai_processes) == 1
    assert ai_processes[0]["pid"] == 4321
    assert ai_processes[0]["is_gpu_process"] is True


def test_get_ai_processes_checks_python_cmdline_without_full_scan(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monitor = AIProcessMonitor()
    non_ai_proc = _FakeProcess(pid=7001, name="python")
    ai_proc = _FakeProcess(pid=7002, name="python")
    cmd_by_pid = {
        7001: "python:app.py",
        7002: "python:train.py",
    }

    monitor._last_full_scan = monitor_module.time.time()

    monkeypatch.setattr(
        monitor_module.psutil, "process_iter", lambda _attrs: [non_ai_proc, ai_proc]
    )
    monkeypatch.setattr(monitor_module.psutil, "pids", lambda: [7001, 7002])
    monkeypatch.setattr(
        monitor, "_get_process_cmdline", lambda proc: cmd_by_pid[proc.pid]
    )
    monkeypatch.setattr(
        monitor,
        "_get_process_info",
        lambda proc, _now, include_extended: {
            "pid": proc.pid,
            "name": proc.info["name"],
            "cmdline": cmd_by_pid[proc.pid],
            "extended": include_extended,
            "cpu_percent": 0.0,
            "memory_percent": 0.0,
            "status": "running",
        },
    )

    ai_processes = monitor.get_ai_processes(gpu_processes=[], force_refresh=False)

    assert [proc["pid"] for proc in ai_processes] == [7002]


def test_get_ai_processes_adds_job_topology_and_straggler_fields(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monitor = AIProcessMonitor()
    rank0 = _FakeProcess(pid=8101, name="python")
    rank1 = _FakeProcess(pid=8102, name="python")

    monitor._last_full_scan = monitor_module.time.time()

    monkeypatch.setattr(
        monitor_module.psutil, "process_iter", lambda _attrs: [rank0, rank1]
    )
    monkeypatch.setattr(monitor_module.psutil, "pids", lambda: [8101, 8102])
    monkeypatch.setattr(
        monitor,
        "_get_process_cmdline",
        lambda proc: (
            "torchrun --rdzv-id run-123 --rank 0 train.py"
            if proc.pid == 8101
            else "torchrun --rdzv-id run-123 --rank 1 train.py"
        ),
    )
    monkeypatch.setattr(monitor, "is_ai_process", lambda _name, _cmdline: True)
    monkeypatch.setattr(
        monitor,
        "_get_process_info",
        lambda proc, _now, include_extended: {
            "pid": proc.pid,
            "name": proc.info["name"],
            "cmdline": (
                "torchrun --rdzv-id run-123 --rank 0 train.py"
                if proc.pid == 8101
                else "torchrun --rdzv-id run-123 --rank 1 train.py"
            ),
            "extended": include_extended,
            "cpu_percent": 80.0 if proc.pid == 8101 else 2.0,
            "memory_percent": 1.0,
            "status": "running",
        },
    )

    ai_processes = monitor.get_ai_processes(
        gpu_processes=[{"pid": 8101, "sm_util": 95.0}, {"pid": 8102, "sm_util": 1.0}],
        force_refresh=False,
    )

    assert len(ai_processes) == 2
    assert {proc["job_id"] for proc in ai_processes} == {"run-run-123"}
    assert {proc["job_size"] for proc in ai_processes} == {2}
    assert {proc["is_distributed_job"] for proc in ai_processes} == {True}

    by_pid = {proc["pid"]: proc for proc in ai_processes}
    assert by_pid[8102]["is_job_straggler"] is True
    assert by_pid[8101]["is_job_straggler"] is False


def test_visibility_grace_window_retains_recent_process_for_120s(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monitor = AIProcessMonitor()
    base_time = 1000.0
    process = {
        "pid": 12001,
        "name": "ollama",
        "create_time": 111.0,
        "cpu_percent": 0.0,
        "memory_percent": 1.0,
        "status": "sleeping",
    }

    initial = monitor._apply_visibility_grace_window([process], base_time)
    assert len(initial) == 1
    assert initial[0]["grace_retained"] is False

    monkeypatch.setattr(monitor, "_is_cached_process_alive", lambda _cached: True)

    retained = monitor._apply_visibility_grace_window([], base_time + 30.0)
    assert len(retained) == 1
    assert retained[0]["pid"] == 12001
    assert retained[0]["grace_retained"] is True
    assert retained[0]["grace_age_s"] == 30.0

    expired = monitor._apply_visibility_grace_window([], base_time + 130.0)
    assert expired == []
