#!/usr/bin/env python3
"""GPU verification of graceful and hard-kill OpProf accounting."""

from __future__ import annotations

import json
import os
import shlex
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_ROOT = WORKDIR / "runs/e-b-sidecar-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")
FLUSH_INTERVAL_SECONDS = 1.0
CHECKPOINT_TOLERANCE_SECONDS = 0.1


def atomic_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f"{path.name}.tmp-{os.getpid()}")
    with temporary.open("w", encoding="utf-8") as output:
        json.dump(value, output, indent=2, sort_keys=True)
        output.write("\n")
        output.flush()
        os.fsync(output.fileno())
    os.replace(temporary, 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(value.strip()) for value 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 assert_idle() -> list[dict[str, int]]:
    rows = gpu_rows()
    if any(row["memory_mib"] or row["utilization_pct"] for row in rows):
        raise RuntimeError(f"dash0 is not GPU-idle: {rows}")
    applications = compute_apps()
    if applications:
        raise RuntimeError(f"compute applications present: {applications}")
    return rows


def wait_ready(process: subprocess.Popen[Any], port: int) -> None:
    deadline = time.monotonic() + 300
    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 wait_zero() -> tuple[list[list[dict[str, int]]], float]:
    samples: list[list[dict[str, int]]] = []
    consecutive = 0
    deadline = time.monotonic() + 180
    while time.monotonic() < deadline and consecutive < 3:
        rows = gpu_rows()
        samples.append(rows)
        zero = all(row["memory_mib"] == 0 for row in rows) and not compute_apps()
        consecutive = consecutive + 1 if zero else 0
        if consecutive < 3:
            time.sleep(2)
    if consecutive < 3:
        raise RuntimeError("GPUs did not reach three stable zero samples")
    return samples, time.time()


def wait_fresh_sidecar(run_dir: Path) -> dict[str, Any]:
    deadline = time.monotonic() + 10
    while time.monotonic() < deadline:
        files = sorted((run_dir / "opprof").glob("*.jsonl.footer.json"))
        if len(files) == 1:
            sidecar = json.loads(files[0].read_text())
            age = time.time_ns() - sidecar["checkpoint_wall_ns"]
            if 0 <= age <= 250_000_000:
                return sidecar
        time.sleep(0.02)
    raise TimeoutError("could not observe a fresh OpProf sidecar")


def validate_accounting(
    run_dir: Path, mode: str, termination_wall_ns: int
) -> dict[str, Any]:
    streams = sorted((run_dir / "opprof").glob("*.jsonl"))
    sidecars = sorted((run_dir / "opprof").glob("*.jsonl.footer.json"))
    if len(streams) != 1 or len(sidecars) != 1:
        raise RuntimeError(
            f"expected one stream/sidecar, got {len(streams)}/{len(sidecars)}"
        )
    raw = streams[0].read_bytes()
    if not raw.endswith(b"\n"):
        raise RuntimeError("partial final JSONL line")
    decoded = [json.loads(line) for line in raw.splitlines()]
    footers = [row for row in decoded if row.get("record_type") == "footer"]
    records = [row for row in decoded if row.get("record_type") != "footer"]
    sidecar = json.loads(sidecars[0].read_text())
    indices = [row["step_index"] for row in records]
    checkpoint_age = (
        termination_wall_ns - sidecar["checkpoint_wall_ns"]
    ) / 1e9
    common = {
        "all_schema_1": all(row.get("schema") == 1 for row in decoded)
        and sidecar.get("schema") == 1,
        "steps_contiguous": indices == list(range(len(indices))),
        "written_matches_records": sidecar["written_records"] == len(records),
        "encoded_balanced": sidecar["encoded_records"]
        == sidecar["written_records"] + sidecar["dropped_records"],
        "last_step_matches": bool(records)
        and sidecar["last_step_index"] == records[-1]["step_index"],
        "zero_drops": sidecar["dropped_records"] == 0
        and all(row["dropped_records_before"] == 0 for row in records),
    }
    if mode == "graceful":
        footer_ok = len(footers) == 1 and decoded[-1] is footers[0]
        agreement = footer_ok and all(
            footers[0][counter] == sidecar[counter]
            for counter in ("encoded_records", "written_records", "dropped_records")
        )
        specific = {
            "one_footer_last": footer_ok,
            "final_sidecar": sidecar["final"] is True,
            "footer_sidecar_agree": agreement,
        }
    else:
        specific = {
            "no_in_stream_footer": not footers,
            "checkpoint_sidecar": sidecar["final"] is False,
            "checkpoint_within_bound": checkpoint_age
            <= FLUSH_INTERVAL_SECONDS + CHECKPOINT_TOLERANCE_SECONDS,
        }
    invariants = {**common, **specific}
    if not all(invariants.values()):
        raise RuntimeError(
            f"{mode} accounting invalid: {invariants}; sidecar={sidecar}"
        )
    return {
        "stream": str(streams[0]),
        "sidecar": str(sidecars[0]),
        "bytes": len(raw),
        "records": len(records),
        "footer_count": len(footers),
        "checkpoint_age_seconds": checkpoint_age,
        "counters": {
            key: sidecar[key]
            for key in ("encoded_records", "written_records", "dropped_records")
        },
        "last_step_index": sidecar["last_step_index"],
        "invariants": invariants,
    }


def run_trial(mode: str, port: int) -> dict[str, Any]:
    run_dir = RUN_ROOT / mode
    if run_dir.exists():
        raise RuntimeError(f"refusing to overwrite {run_dir}")
    run_dir.mkdir(parents=True)
    before = assert_idle()
    atomic_json(run_dir / "gpu-before.json", before)
    (run_dir / "clocks-before.txt").write_text(
        run_text(["nvidia-smi", "-q", "-d", "CLOCK"])
    )
    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",
        "--shutdown-timeout",
        "120",
    ]
    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",
        "20",
        "--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", encoding="utf-8") as output:
        output.write(
            f"GPU_COMMAND {mode} server: {shlex.join(server_command)}; "
            "expected=60-180s startup + 140-180s load\n"
        )
        output.write(
            f"GPU_COMMAND {mode} client: {shlex.join(client_command)}; "
            "expected=20s warmup + 120s clean + drain\n"
        )
    environment = os.environ.copy()
    environment.update(
        {
            "CUDA_VISIBLE_DEVICES": "0",
            "VLLM_OPPROF_DIR": str(run_dir / "opprof"),
            "HF_HUB_OFFLINE": "1",
            "TRANSFORMERS_OFFLINE": "1",
            "PYTHONUNBUFFERED": "1",
        }
    )
    started = time.time()
    server_log = (run_dir / "server.log").open("ab", buffering=0)
    server: subprocess.Popen[Any] | None = None
    try:
        print(
            f"GPU_COMMAND sidecar-{mode}: P01/C00 GPU0, 20s warmup + "
            "120s clean, expected 4-6 wall-min",
            flush=True,
        )
        server = subprocess.Popen(
            server_command,
            cwd=SOURCE,
            env=environment,
            stdout=server_log,
            stderr=subprocess.STDOUT,
            start_new_session=True,
        )
        atomic_json(
            run_dir / "state.json",
            {"status": "server_starting", "server_pid": server.pid},
        )
        wait_ready(server, port)
        with (run_dir / "client.log").open("ab", buffering=0) as client_log:
            client = subprocess.run(
                client_command,
                cwd=WORKDIR,
                stdout=client_log,
                stderr=subprocess.STDOUT,
            )
        if client.returncode:
            raise RuntimeError(f"client exited {client.returncode}")
        if mode == "graceful":
            termination_wall_ns = time.time_ns()
            os.kill(server.pid, signal.SIGINT)
            server.wait(timeout=150)
            log_text = (run_dir / "server.log").read_text(errors="replace")
            if "mode=drain timeout=120s" not in log_text:
                raise RuntimeError("official drain-mode log not observed")
        else:
            wait_fresh_sidecar(run_dir)
            termination_wall_ns = time.time_ns()
            os.killpg(server.pid, signal.SIGKILL)
            server.wait(timeout=30)
        zero_samples, zero_at = wait_zero()
        accounting = validate_accounting(run_dir, mode, termination_wall_ns)
        client_result = json.loads((run_dir / "client/result.json").read_text())
        result = {
            "schema": 1,
            "status": "pass",
            "mode": mode,
            "server_returncode": server.returncode,
            "termination_wall_ns": termination_wall_ns,
            "accounting": accounting,
            "clean": client_result["clean"],
            "failed_records": client_result["failed_records"],
            "gpu_zero_samples": zero_samples,
            "gpu_seconds": zero_at - started,
        }
        atomic_json(run_dir / "result.json", result)
        atomic_json(run_dir / "state.json", {"status": "complete"})
        return result
    except Exception as error:
        atomic_json(
            run_dir / "state.json",
            {"status": "failed", "failure": repr(error)},
        )
        if server is not None and server.poll() is None:
            try:
                os.killpg(server.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
        wait_zero()
        raise
    finally:
        server_log.close()


def main() -> None:
    if RUN_ROOT.exists():
        raise RuntimeError(f"refusing to overwrite {RUN_ROOT}")
    RUN_ROOT.mkdir(parents=True)
    state: dict[str, Any] = {"schema": 1, "status": "running", "results": {}}
    atomic_json(RUN_ROOT / "state.json", state)
    for offset, mode in enumerate(("graceful", "hard-kill")):
        result = run_trial(mode, 8010 + offset)
        state["results"][mode] = result
        atomic_json(RUN_ROOT / "state.json", state)
    state["status"] = "complete"
    state["gpu_seconds"] = sum(
        result["gpu_seconds"] for result in state["results"].values()
    )
    atomic_json(RUN_ROOT / "state.json", state)
    print(json.dumps(state, sort_keys=True), flush=True)


if __name__ == "__main__":
    main()