"""Tests for the GPU process list + kill endpoints.

Standalone (no pytest required):

    python tests/test_processes.py

Also works under pytest if available. Covers the /proc reader (including
comm values with spaces and parentheses), the kill endpoint's input
validation, a real SIGTERM round-trip, and the GET /api/processes payload.
"""

import contextlib
import os
import subprocess
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from nvcurve import server  # noqa: E402
from nvcurve.hal import processes  # noqa: E402

PASS = 0
FAIL = 0


def check(name: str, cond: bool) -> None:
    global PASS, FAIL
    if cond:
        PASS += 1
        print(f"  PASS  {name}")
    else:
        FAIL += 1
        print(f"  FAIL  {name}")


def test_proc_info_self() -> None:
    """_proc_info must return sane data for our own process."""
    info = processes._proc_info(os.getpid())
    check("self info not None", info is not None)
    if info is None:
        return
    check(
        "user is non-empty string",
        bool(isinstance(info["user"], str) and info["user"]),
    )
    check(
        "cpu_pct is None or non-negative number",
        info["cpu_pct"] is None or (isinstance(info["cpu_pct"], float) and info["cpu_pct"] >= 0),
    )
    check("mem_bytes > 0", info["mem_bytes"] is not None and info["mem_bytes"] > 0)
    check("command non-empty", isinstance(info["command"], str) and len(info["command"]) > 0)


def test_proc_info_missing() -> None:
    """A pid that does not exist must yield None (not an exception)."""
    check("missing pid -> None", processes._proc_info(99999999) is None)


def test_proc_info_comm_with_parens() -> None:
    """comm containing spaces and parentheses must not break /proc/stat parsing.

    The kernel comm is set via prctl(PR_SET_NAME); `exec -a` only changes
    argv[0] and would leave comm as the executable name, so the parens path
    would never be exercised.
    """
    proc = subprocess.Popen(
        [
            sys.executable,
            "-c",
            "import ctypes, time\n"
            "ctypes.CDLL(None).prctl(15, b'proc (test) x', 0, 0, 0)\n"
            "time.sleep(30)",
        ]
    )
    try:
        # Wait until /proc/<pid>/stat exists and comm actually carries the
        # parenthesised name (prctl may take a moment after exec).
        comm = ""
        for _ in range(100):
            try:
                with open(f"/proc/{proc.pid}/stat") as f:
                    raw = f.read()
                comm = raw[raw.index("(") + 1 : raw.rfind(")")]
            except OSError:
                pass
            if "(" in comm:
                break
            time.sleep(0.05)
        check("fixture comm contains parens", "(" in comm)
        info = processes._proc_info(proc.pid)
        check("weird-comm info not None", info is not None)
        if info is not None:
            check("weird-comm command parsed", "prctl" in info["command"])
    finally:
        proc.kill()
        proc.wait()


def test_kill_endpoint_validation() -> None:
    """The kill endpoint must reject bad input and unknown pids."""
    from fastapi.testclient import TestClient

    client = TestClient(server.app)
    r = client.post(
        "/api/processes/kill", json={"pid": os.getpid(), "signal": "BOGUS"}
    )
    check("invalid signal -> 400", r.status_code == 400)
    r = client.post("/api/processes/kill", json={"pid": -5, "signal": "TERM"})
    check("negative pid -> 400", r.status_code == 400)
    r = client.post("/api/processes/kill", json={"pid": 0, "signal": "TERM"})
    check("zero pid -> 400", r.status_code == 400)
    r = client.post("/api/processes/kill", json={"pid": 99999999, "signal": "TERM"})
    check("unknown pid -> 404", r.status_code == 404)


def test_kill_endpoint_real() -> None:
    """A real SIGTERM round-trip: the process must actually exit."""
    from fastapi.testclient import TestClient

    client = TestClient(server.app)
    proc = subprocess.Popen(["sleep", "30"])
    try:
        r = client.post(
            "/api/processes/kill", json={"pid": proc.pid, "signal": "TERM"}
        )
        check("kill -> 200", r.status_code == 200)
        check("response ok", r.json().get("ok") is True)
        try:
            proc.wait(timeout=5)
            check("process exited", True)
        except subprocess.TimeoutExpired:
            check("process exited", False)
    finally:
        proc.kill()
        proc.wait()


def test_kill_endpoint_refuses_init_kthreadd() -> None:
    """The endpoint must refuse to signal init (pid 1) or kthreadd (pid 2).

    SIGTERM to init is a system-shutdown vector; kernel threads don't handle
    signals anyway. This fires before any NVML/state access, so no GPU needed.
    """
    from fastapi.testclient import TestClient

    client = TestClient(server.app)
    for pid in (1, 2):
        r = client.post(
            "/api/processes/kill", json={"pid": pid, "signal": "KILL"}
        )
        check(f"pid {pid} refused -> 400", r.status_code == 400)


def test_get_parent_pid() -> None:
    """get_parent_pid must return the kernel ppid and degrade gracefully."""
    check("self ppid matches", processes.get_parent_pid(os.getpid()) == os.getppid())
    check("missing pid -> None", processes.get_parent_pid(99999999) is None)
    check("pid 1 has no parent", processes.get_parent_pid(1) == 0)


def test_kill_parent_endpoint() -> None:
    """parent=true must signal the parent, leaving the child alive."""
    from fastapi.testclient import TestClient

    client = TestClient(server.app)

    # Validation cases first (no fixture needed).
    r = client.post(
        "/api/processes/kill", json={"pid": 99999999, "signal": "KILL", "parent": True}
    )
    check("parent of missing pid -> 404", r.status_code == 404)
    r = client.post(
        "/api/processes/kill", json={"pid": 1, "signal": "KILL", "parent": True}
    )
    check("pid 1 has no parent -> 400", r.status_code == 400)

    # bash forks a background sleep (child), then sleeps itself (parent).
    parent = subprocess.Popen(["bash", "-c", "sleep 30 & sleep 30"])
    child: int | None = None
    for _ in range(100):
        try:
            with open(f"/proc/{parent.pid}/task/{parent.pid}/children") as f:
                kids = f.read().split()
            if kids:
                child = int(kids[0])
                break
        except OSError:
            pass
        time.sleep(0.05)
    try:
        check("fixture has a child", child is not None)
        if child is None:
            return
        r = client.post(
            "/api/processes/kill",
            json={"pid": child, "signal": "KILL", "parent": True},
        )
        check("kill parent -> 200", r.status_code == 200)
        check("response targets the parent", r.json().get("pid") == parent.pid)
        check("response reports parent flag", r.json().get("parent") is True)
        try:
            parent.wait(timeout=5)
            check("parent died", True)
        except subprocess.TimeoutExpired:
            check("parent died", False)
        time.sleep(0.2)
        try:
            os.kill(child, 0)
            check("child survived (orphaned)", True)
        except ProcessLookupError:
            check("child survived (orphaned)", False)
    finally:
        for victim in (parent,):
            with contextlib.suppress(ProcessLookupError, subprocess.TimeoutExpired):
                victim.kill()
                victim.wait(timeout=5)
        if child is not None:
            with contextlib.suppress(ProcessLookupError):
                os.kill(child, 9)


def test_processes_endpoint() -> None:
    """GET /api/processes must return the HAL payload for a known GPU."""
    from fastapi.testclient import TestClient

    fake = {
        "processes": [
            {
                "pid": 1,
                "user": "root",
                "dev": 3,
                "type": "C",
                "gpu_util_pct": 42.0,
                "mem_util_pct": 1.0,
                "vram_bytes": 1000,
                "vram_pct": 0.1,
                "cpu_pct": 1.5,
                "mem_host_bytes": 2048,
                "command": "/usr/bin/fake",
            }
        ],
        "mem_total_bytes": 1000000,
    }
    orig = server.list_gpu_processes
    server._state["gpus"][3] = {"gpu": object(), "gpu_name": "fake"}
    server.list_gpu_processes = lambda idx: fake
    try:
        client = TestClient(server.app)
        r = client.get("/api/processes", params={"gpu_index": 3})
        check("processes -> 200", r.status_code == 200)
        check("payload matches HAL", r.json() == fake)
        r = client.get("/api/processes", params={"gpu_index": 99})
        check("unknown gpu -> 404", r.status_code == 404)
    finally:
        server.list_gpu_processes = orig
        server._state["gpus"].pop(3, None)


def test_list_gpu_processes_shape() -> None:
    """With NVML available, list_gpu_processes must return a well-shaped payload."""
    from nvcurve.hal import monitoring

    if not monitoring._NVML_AVAILABLE:
        print("  SKIP  NVML not available")
        return
    if not monitoring.init_nvml():
        print("  SKIP  NVML init failed (no driver?)")
        return
    try:
        data = processes.list_gpu_processes(0)
        check("has processes list", isinstance(data.get("processes"), list))
        check("has mem_total_bytes", "mem_total_bytes" in data)
        keys = {
            "pid",
            "user",
            "dev",
            "type",
            "gpu_util_pct",
            "mem_util_pct",
            "vram_bytes",
            "vram_pct",
            "cpu_pct",
            "mem_host_bytes",
            "command",
        }
        for p in data["processes"]:
            if not keys.issubset(p.keys()):
                check(f"process {p.get('pid')} has all keys", False)
                break
        else:
            check("all processes have all keys", True)
        # Sorted by VRAM descending.
        vram = [p["vram_bytes"] or 0 for p in data["processes"]]
        check("sorted by VRAM desc", vram == sorted(vram, reverse=True))
    finally:
        monitoring.shutdown_nvml()


def test_list_gpu_processes_no_nvml() -> None:
    """With NVML unavailable, list_gpu_processes returns an empty payload.

    This is the branch the shape test skips (no driver): the UI uses
    mem_total_bytes is None to render 'data unavailable' instead of
    'no processes'.
    """
    from nvcurve.hal import monitoring

    orig_avail = monitoring._NVML_AVAILABLE
    orig_init = monitoring._nvml_initialized
    monitoring._NVML_AVAILABLE = False
    monitoring._nvml_initialized = False
    try:
        data = processes.list_gpu_processes(0)
        check("empty processes list", data["processes"] == [])
        check("mem_total_bytes is None", data["mem_total_bytes"] is None)
    finally:
        monitoring._NVML_AVAILABLE = orig_avail
        monitoring._nvml_initialized = orig_init


def main() -> None:
    tests = [
        test_proc_info_self,
        test_proc_info_missing,
        test_proc_info_comm_with_parens,
        test_kill_endpoint_validation,
        test_kill_endpoint_real,
        test_kill_endpoint_refuses_init_kthreadd,
        test_get_parent_pid,
        test_kill_parent_endpoint,
        test_processes_endpoint,
        test_list_gpu_processes_shape,
        test_list_gpu_processes_no_nvml,
    ]
    for t in tests:
        print(f"== {t.__name__} ==")
        t()
    print()
    print(f"{PASS} passed, {FAIL} failed")
    if FAIL:
        sys.exit(1)


if __name__ == "__main__":
    main()