#!/usr/bin/env python3
"""One-GPU verification of API-parent-first OpProf shutdown."""

from __future__ import annotations

import json
import os
import shlex
import shutil
import signal
import subprocess
import time
import urllib.request
from pathlib import Path
from typing import Any

WORKDIR = Path("/home/admin/cpfs/wjh/opprof-phase3-dash0-20260712")
RUN_DIR = WORKDIR / "runs/e-b-shutdown-verification"
SOURCE = Path(
    "/home/admin/cpfs/wjh/opprof-phase2-dash0-20260711/vllm-v0.24.0"
)
VENV = Path("/tmp/wjh-opprof-phase2-dash0-20260711/.venv")
MODEL = Path("/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B")
CLIENT = WORKDIR / "scripts/opprof_phase3_client.py"
MANIFEST = Path("/home/admin/cpfs/wjh/opprof-phase3-private/manifests/P01.jsonl")
PORT = 8010


def atomic_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(path.name + f".tmp.{os.getpid()}")
    with tmp.open("w", encoding="utf-8") as f:
        json.dump(value, f, sort_keys=True, indent=2)
        f.write("\n")
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)


def run_text(command: list[str], check: bool = True) -> str:
    result = subprocess.run(
        command, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
    )
    if check and result.returncode:
        raise RuntimeError(
            f"command failed ({result.returncode}): {shlex.join(command)}\n"
            f"{result.stdout}"
        )
    return result.stdout


def gpu_rows() -> list[dict[str, int]]:
    output = run_text(
        [
            "nvidia-smi",
            "--query-gpu=index,memory.used,utilization.gpu",
            "--format=csv,noheader,nounits",
        ]
    )
    rows = []
    for line in output.strip().splitlines():
        index, memory, utilization = [int(part.strip()) for part in line.split(",")]
        rows.append(
            {"index": index, "memory_mib": memory, "utilization_pct": utilization}
        )
    return rows


def compute_apps() -> str:
    return run_text(
        [
            "nvidia-smi",
            "--query-compute-apps=gpu_uuid,pid,process_name,used_memory",
            "--format=csv,noheader,nounits",
        ],
        check=False,
    ).strip()


def group_alive(pgid: int, process: subprocess.Popen[Any] | None = None) -> bool:
    if process is not None:
        process.poll()  # reap an exited API parent before probing its group
    try:
        os.killpg(pgid, 0)
        return True
    except ProcessLookupError:
        return False


def wait_ready(process: subprocess.Popen[Any], timeout: float = 300) -> None:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if process.poll() is not None:
            raise RuntimeError("server exited before readiness")
        try:
            with urllib.request.urlopen(
                f"http://127.0.0.1:{PORT}/health", timeout=1
            ) as response:
                if response.status == 200:
                    return
        except Exception:
            pass
        time.sleep(1)
    raise TimeoutError("server readiness timeout")


def shutdown_parent_first(process: subprocess.Popen[Any]) -> dict[str, Any]:
    called_at = time.time()
    os.kill(process.pid, signal.SIGINT)
    deadline = time.monotonic() + 120
    while time.monotonic() < deadline and group_alive(process.pid, process):
        time.sleep(1)
    graceful = not group_alive(process.pid, process)
    if not graceful:
        for sig in (signal.SIGTERM, signal.SIGKILL):
            try:
                os.killpg(process.pid, sig)
            except ProcessLookupError:
                break
            time.sleep(5)
    return {
        "api_parent_pid": process.pid,
        "sigint_wall_time": called_at,
        "group_gone_wall_time": time.time(),
        "graceful": graceful,
        "returncode": process.poll(),
    }


def wait_stable_zero() -> list[list[dict[str, int]]]:
    samples = []
    consecutive = 0
    deadline = time.monotonic() + 120
    while time.monotonic() < deadline and consecutive < 3:
        rows = gpu_rows()
        samples.append(rows)
        zero = rows[0]["memory_mib"] == 0 and not compute_apps()
        consecutive = consecutive + 1 if zero else 0
        if consecutive < 3:
            time.sleep(5)
    if consecutive < 3:
        raise RuntimeError("GPU0 did not reach three stable zero samples")
    return samples


def validate_footer() -> dict[str, Any]:
    files = sorted((RUN_DIR / "opprof").glob("*.jsonl"))
    if len(files) != 1:
        raise RuntimeError(f"expected one OpProf JSONL, found {len(files)}")
    lines = files[0].read_text().splitlines()
    decoded = [json.loads(line) for line in lines]
    if not decoded or decoded[-1].get("record_type") != "footer":
        raise RuntimeError("Layer-1 footer missing")
    records, footer = decoded[:-1], decoded[-1]
    indices = [record["step_index"] for record in records]
    invariants = {
        "all_schema_1": all(item.get("schema") == 1 for item in decoded),
        "one_footer_last": sum(
            item.get("record_type") == "footer" for item in decoded
        )
        == 1,
        "steps_contiguous": indices == list(range(len(indices))),
        "written_matches_records": footer["written_records"] == len(records),
        "encoded_balanced": footer["encoded_records"]
        == footer["written_records"] + footer["dropped_records"],
        "zero_drops": footer["dropped_records"] == 0
        and all(record["dropped_records_before"] == 0 for record in records),
    }
    if not all(invariants.values()):
        raise RuntimeError(f"footer accounting invalid: {invariants}, {footer}")
    return {
        "path": str(files[0]),
        "bytes": files[0].stat().st_size,
        "records": len(records),
        "footer": footer,
        "invariants": invariants,
    }


def main() -> None:
    if RUN_DIR.exists():
        raise RuntimeError(f"refusing to overwrite {RUN_DIR}")
    RUN_DIR.mkdir(parents=True)
    state: dict[str, Any] = {
        "schema": 1,
        "status": "preflight",
        "started_at": time.time(),
        "gpu": 0,
    }
    atomic_json(RUN_DIR / "state.json", state)
    before = gpu_rows()
    if before[0]["memory_mib"] != 0 or before[0]["utilization_pct"] != 0:
        raise RuntimeError(f"GPU0 is not idle: {before[0]}")
    if compute_apps():
        raise RuntimeError("a compute process exists before verification")
    (RUN_DIR / "gpu-before.json").write_text(json.dumps(before, indent=2) + "\n")
    (RUN_DIR / "clocks-before.txt").write_text(
        run_text(["nvidia-smi", "-q", "-d", "CLOCK"])
    )
    profile_config = {
        "profiler": "torch",
        "torch_profiler_dir": "/tmp/wjh-opprof-p3-footer-verify",
        "ignore_frontend": True,
        "wait_iterations": 0,
        "warmup_iterations": 2,
        "active_iterations": 8,
    }
    trace_dir = Path(profile_config["torch_profiler_dir"])
    if trace_dir.exists():
        shutil.rmtree(trace_dir)
    trace_dir.mkdir()
    server_command = [
        "taskset",
        "-c",
        "0-19",
        str(VENV / "bin/vllm"),
        "serve",
        str(MODEL),
        "--host",
        "127.0.0.1",
        "--port",
        str(PORT),
        "--tensor-parallel-size",
        "1",
        "--enable-chunked-prefill",
        "--enable-prefix-caching",
        "--profiler-config",
        json.dumps(profile_config, separators=(",", ":")),
    ]
    client_command = [
        "taskset",
        "-c",
        "0-19",
        str(VENV / "bin/python"),
        str(CLIENT),
        "run",
        "--manifest",
        str(MANIFEST),
        "--base-url",
        f"http://127.0.0.1:{PORT}",
        "--model",
        str(MODEL),
        "--load-point",
        "saturation",
        "--request-rate",
        "inf",
        "--max-concurrency",
        "256",
        "--ignore-eos",
        "--temperature",
        "0",
        "--warmup-seconds",
        "30",
        "--clean-segment-seconds",
        "40",
        "--num-clean-segments",
        "3",
        "--drain-timeout-seconds",
        "120",
        "--workload-seed",
        "20260712",
        "--result-dir",
        str(RUN_DIR / "client"),
    ]
    with (RUN_DIR / "commands.log").open("w") as f:
        f.write(
            "GPU_COMMAND shutdown footer server: "
            + shlex.join(server_command)
            + " ; expected=60-180s startup + 150-180s load\n"
        )
        f.write(
            "GPU_COMMAND shutdown footer client: "
            + shlex.join(client_command)
            + " ; expected=30s warmup + 120s clean + drain\n"
        )
    server_log = (RUN_DIR / "server.log").open("ab", buffering=0)
    env = os.environ.copy()
    env.update(
        {
            "CUDA_VISIBLE_DEVICES": "0",
            "VLLM_OPPROF_DIR": str(RUN_DIR / "opprof"),
            "HF_HUB_OFFLINE": "1",
            "TRANSFORMERS_OFFLINE": "1",
            "PYTHONUNBUFFERED": "1",
        }
    )
    server: subprocess.Popen[Any] | None = None
    try:
        print(
            "GPU_COMMAND shutdown-fix verification: P01/C00 GPU0, "
            "30s warmup + 120s clean, expected 4-6 wall-min",
            flush=True,
        )
        server = subprocess.Popen(
            server_command,
            cwd=SOURCE,
            env=env,
            stdout=server_log,
            stderr=subprocess.STDOUT,
            start_new_session=True,
        )
        state.update({"status": "server_starting", "server_pid": server.pid})
        atomic_json(RUN_DIR / "state.json", state)
        wait_ready(server)
        state.update({"status": "client_running", "server_ready_at": time.time()})
        atomic_json(RUN_DIR / "state.json", state)
        with (RUN_DIR / "client.log").open("ab", buffering=0) as client_log:
            result = subprocess.run(
                client_command,
                cwd=WORKDIR,
                stdout=client_log,
                stderr=subprocess.STDOUT,
            )
        if result.returncode:
            raise RuntimeError(f"client exited {result.returncode}")
        state["status"] = "parent_first_shutdown"
        atomic_json(RUN_DIR / "state.json", state)
        shutdown = shutdown_parent_first(server)
        if not shutdown["graceful"]:
            raise RuntimeError(f"server required escalation: {shutdown}")
        footer = validate_footer()
        zero_samples = wait_stable_zero()
        client_result = json.loads((RUN_DIR / "client/result.json").read_text())
        result_json = {
            "schema": 1,
            "status": "pass",
            "gpu": 0,
            "shutdown": shutdown,
            "footer": footer,
            "clean": client_result["clean"],
            "failed_records": client_result["failed_records"],
            "gpu_zero_samples": zero_samples,
            "gpu_seconds": zero_samples and time.time() - state["started_at"],
        }
        atomic_json(RUN_DIR / "result.json", result_json)
        state.update({"status": "complete", "completed_at": time.time()})
        atomic_json(RUN_DIR / "state.json", state)
        print(json.dumps(result_json, sort_keys=True), flush=True)
    except Exception as error:
        state.update({"status": "failed", "failure": repr(error)})
        atomic_json(RUN_DIR / "state.json", state)
        if server is not None and group_alive(server.pid, server):
            try:
                os.killpg(server.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
        raise
    finally:
        server_log.close()


if __name__ == "__main__":
    main()