#!/usr/bin/env python3
"""One-pair P09 CUDAGraph capture-size validation for OpProf Phase 4."""

from __future__ import annotations

import argparse
import json
import math
import os
import shlex
import signal
import subprocess
import sys
import time
import urllib.request
from pathlib import Path
from typing import Any

import numpy as np

import opprof_phase3_controller as common
import opprof_phase3_matrix as matrix


WORKDIR = Path("/home/admin/cpfs/wjh/opprof-phase3-dash0-20260712")
ROOT = WORKDIR / "runs/phase4-capture-p09"
PHASE3 = WORKDIR / "runs/phase3"
PRIVATE = Path("/home/admin/cpfs/wjh/opprof-phase3-private/manifests")
MODEL = Path("/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B")
SOURCE = Path("/home/admin/cpfs/wjh/opprof-phase2-dash0-20260711/vllm-v0.24.0")
VENV = Path("/tmp/wjh-opprof-phase2-dash0-20260711/.venv")
CLIENT = WORKDIR / "scripts/opprof_phase3_client.py"
STATE = ROOT / "controller-state.json"
GPU = 0
PORT = 8200
PRIOR_GPU_HOURS = 14.025875418755744
GPU_HOUR_LIMIT = 16.0
EXPECTED_INCREMENT_HOURS = 0.5
ADDED_CAPTURE_SIZES = (3, 5, 6, 7, 9)
DEFAULT_CAPTURE_SIZES = (
    1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104,
    112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200,
    208, 216, 224, 232, 240, 248, 256, 272, 288, 304, 320, 336,
    352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512,
)
ON_CAPTURE_SIZES = tuple(sorted(set(DEFAULT_CAPTURE_SIZES + ADDED_CAPTURE_SIZES)))


def plan() -> dict[str, Any]:
    return {
        "schema": 1,
        "pattern": "P09",
        "load": "moderate",
        "order": ["ON", "OFF"],
        "gpu": GPU,
        "warmup_seconds": 60,
        "clean_seconds": 240,
        "profile": False,
        "manifest": str(PRIVATE / "P09.jsonl"),
        "saturation_result": str(
            PHASE3 / "primary/P09-C00/saturation/client/result.json"
        ),
        "rate_fraction": 0.60,
        "added_capture_sizes": list(ADDED_CAPTURE_SIZES),
        "on_capture_sizes": list(ON_CAPTURE_SIZES),
        "measured_padding_recovery_bound": 0.049776380388728225,
        "primary_metric": "clean graph-hit padding_fraction",
        "secondary_metrics": [
            "Layer-1 useful scheduled tokens/model-step millisecond",
            "clean completed request throughput",
            "clean request latency",
        ],
        "prior_gpu_hours": PRIOR_GPU_HOURS,
        "projected_increment_gpu_hours": EXPECTED_INCREMENT_HOURS,
        "projected_total_gpu_hours": PRIOR_GPU_HOURS + EXPECTED_INCREMENT_HOURS,
        "gpu_hour_limit": GPU_HOUR_LIMIT,
    }


def save_state(state: dict[str, Any]) -> None:
    state["updated_at"] = time.time()
    state["controller_pid"] = os.getpid()
    common.atomic_json(STATE, state)


def wait_ready(server: subprocess.Popen[Any]) -> None:
    deadline = time.monotonic() + 300
    while time.monotonic() < deadline:
        if server.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 server_command(arm: str) -> list[str]:
    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",
    ]
    if arm == "ON":
        command.extend(("--cudagraph-capture-sizes", *map(str, ON_CAPTURE_SIZES)))
    return command


def client_command(run_dir: Path) -> list[str]:
    return [
        "taskset", "-c", "0-19", str(VENV / "bin/python"), str(CLIENT), "run",
        "--manifest", str(PRIVATE / "P09.jsonl"),
        "--base-url", f"http://127.0.0.1:{PORT}", "--model", str(MODEL),
        "--load-point", "moderate", "--saturation-result",
        str(PHASE3 / "primary/P09-C00/saturation/client/result.json"),
        "--rate-fraction", "0.60", "--max-concurrency", "256", "--ignore-eos",
        "--temperature", "0", "--warmup-seconds", "60",
        "--clean-segment-seconds", "80", "--num-clean-segments", "3",
        "--recovery-seconds", "30", "--drain-timeout-seconds", "120",
        "--workload-seed", "20260712", "--server-seed", "20260712",
        "--result-dir", str(run_dir / "client"),
    ]


def summarize(run_dir: Path, arm: str) -> dict[str, Any]:
    result = json.loads((run_dir / "client/result.json").read_text())
    requests = [
        json.loads(line)
        for line in (run_dir / "client/requests.jsonl").read_text().splitlines()
    ]
    t0 = int(result["t0_mono_ns"])
    start, end = t0 + int(60e9), t0 + int(300e9)
    stream = next((run_dir / "opprof").glob("*.jsonl"))
    records = []
    for line in stream.read_text().splitlines():
        record = json.loads(line)
        if (
            "step_index" in record
            and start <= int(record["submit_mono_ns"]) < end
        ):
            records.append(record)
    model = [record for record in records if record["model_executed"]]
    hits = [
        record
        for record in model
        if record["cudagraph"]["hit"]
        and int(record["cudagraph"]["bucket_tokens"]) > 0
    ]
    pad = sum(int(record["cudagraph"]["padding_tokens"]) for record in hits)
    bucket = sum(int(record["cudagraph"]["bucket_tokens"]) for record in hits)
    useful = sum(
        int(record["prefill_tokens"]) + int(record["decode_tokens"])
        for record in records
    )
    duration_ms = sum(
        (int(record["complete_mono_ns"]) - int(record["submit_mono_ns"])) / 1e6
        for record in records
    )
    completed = [
        request
        for request in requests
        if request["success"] and 60 <= float(request["completed_s"]) < 300
    ]
    e2e = np.asarray(
        [float(request["completed_s"] - request["admitted_s"]) for request in completed]
    )
    layer1 = matrix.validate_layer1(run_dir)
    return {
        "schema": 1,
        "arm": arm,
        "capture_sizes": list(ON_CAPTURE_SIZES) if arm == "ON" else "default",
        "clean_completed": len(completed),
        "clean_failed": int(result["clean"]["failed"]),
        "clean_completed_throughput_rps": float(
            result["clean"]["completed_throughput_rps"]
        ),
        "clean_offered_rps": float(result["clean"]["offered_rps"]),
        "e2e_latency_mean_s": float(e2e.mean()),
        "e2e_latency_p95_s": float(np.quantile(e2e, 0.95)),
        "model_steps": len(model),
        "graph_hit_steps": len(hits),
        "padding_tokens": pad,
        "bucket_tokens": bucket,
        "padding_fraction": pad / bucket,
        "graph_miss_rate": sum(not record["cudagraph"]["hit"] for record in model)
        / len(model),
        "useful_tokens": useful,
        "model_step_duration_ms": duration_ms,
        "token_efficiency_per_ms": useful / duration_ms,
        "layer1_records": layer1["records"],
        "layer1_invariants": layer1["invariants"],
        "drain_seconds": float(result["drain_seconds"]),
    }


def run_arm(state: dict[str, Any], arm: str) -> None:
    if state["arms"].get(arm, {}).get("status") == "complete":
        return
    run_dir = ROOT / arm.lower()
    if run_dir.exists():
        run_dir.rename(run_dir.with_name(f"{run_dir.name}.interrupted-{int(time.time())}"))
    run_dir.mkdir(parents=True)
    common.preflight([GPU], run_dir)
    server_cmd = server_command(arm)
    client_cmd = client_command(run_dir)
    commands_path = run_dir / "commands.log"
    common.command_log(commands_path, f"P09-capture-{arm}-server", server_cmd, "6-9m")
    common.command_log(commands_path, f"P09-capture-{arm}-client", client_cmd, "5-7m")
    print(
        f"GPU_COMMAND P09-capture-{arm}-server: {shlex.join(server_cmd)}; "
        "expected=6-9m",
        flush=True,
    )
    print(
        f"GPU_COMMAND P09-capture-{arm}-client: {shlex.join(client_cmd)}; "
        "expected=5-7m",
        flush=True,
    )
    environment = os.environ.copy()
    environment.update(
        {
            "CUDA_VISIBLE_DEVICES": str(GPU),
            "VLLM_OPPROF_DIR": str(run_dir / "opprof"),
            "HF_HUB_OFFLINE": "1",
            "TRANSFORMERS_OFFLINE": "1",
            "PYTHONUNBUFFERED": "1",
        }
    )
    server_log = (run_dir / "server.log").open("ab", buffering=0)
    server_started = time.time()
    server = subprocess.Popen(
        server_cmd,
        cwd=SOURCE,
        env=environment,
        stdout=server_log,
        stderr=subprocess.STDOUT,
        start_new_session=True,
    )
    client = None
    client_log = None
    monitor = None
    owned = {server.pid}
    state["arms"][arm] = {
        "status": "starting",
        "server_pid": server.pid,
        "started_at": server_started,
    }
    save_state(state)
    failure = None
    try:
        wait_ready(server)
        monitor = common.Monitor(run_dir / "monitor.jsonl", owned)
        monitor.start()
        client_log = (run_dir / "client.log").open("ab", buffering=0)
        client = subprocess.Popen(
            client_cmd,
            cwd=WORKDIR,
            stdout=client_log,
            stderr=subprocess.STDOUT,
            start_new_session=True,
        )
        owned.add(client.pid)
        state["arms"][arm].update(status="running", client_pid=client.pid)
        save_state(state)
        deadline = time.monotonic() + 900
        while client.poll() is None and time.monotonic() < deadline:
            if server.poll() is not None:
                raise RuntimeError("server exited during client load")
            if monitor.other_apps:
                raise RuntimeError(f"other GPU process appeared: {monitor.other_apps}")
            time.sleep(2)
        if client.poll() is None:
            raise TimeoutError("client exceeded 900 seconds")
        if client.returncode:
            raise RuntimeError(f"client failed with {client.returncode}")
    except Exception as error:
        failure = error
    finally:
        if client is not None and client.poll() is None:
            try:
                os.killpg(client.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
        common.stop_servers([server])
        server_log.close()
        if client_log is not None:
            client_log.close()
        if monitor is not None:
            monitor.stop()
        common.verify_idle([GPU], run_dir)
    gpu_hours = (time.time() - server_started) / 3600
    state["gpu_hours_increment"] += gpu_hours
    if PRIOR_GPU_HOURS + state["gpu_hours_increment"] >= GPU_HOUR_LIMIT:
        failure = failure or RuntimeError("GPU-hour limit reached")
    if failure is not None:
        state["arms"][arm].update(status="failed", failure=repr(failure))
        state["status"] = "failed"
        save_state(state)
        raise failure
    summary = summarize(run_dir, arm)
    summary["gpu_hours"] = gpu_hours
    common.atomic_json(run_dir / "summary.json", summary)
    state["arms"][arm].update(status="complete", summary=summary)
    save_state(state)


def run() -> None:
    ROOT.mkdir(parents=True, exist_ok=True)
    if PRIOR_GPU_HOURS + EXPECTED_INCREMENT_HOURS >= GPU_HOUR_LIMIT:
        raise RuntimeError("projected validation exceeds GPU-hour budget")
    if STATE.exists():
        state = json.loads(STATE.read_text())
    else:
        state = {
            "schema": 1,
            "status": "running",
            "created_at": time.time(),
            "plan": plan(),
            "arms": {},
            "gpu_hours_increment": 0.0,
        }
    for arm in ("ON", "OFF"):
        run_arm(state, arm)
    on, off = state["arms"]["ON"]["summary"], state["arms"]["OFF"]["summary"]
    result = {
        "schema": 1,
        "plan": plan(),
        "on": on,
        "off": off,
        "delta": {
            "padding_fraction_points": on["padding_fraction"] - off["padding_fraction"],
            "padding_reduction_fraction": 1 - on["padding_fraction"] / off["padding_fraction"],
            "token_efficiency_relative": on["token_efficiency_per_ms"]
            / off["token_efficiency_per_ms"]
            - 1,
            "completed_throughput_relative": on["clean_completed_throughput_rps"]
            / off["clean_completed_throughput_rps"]
            - 1,
            "e2e_mean_latency_relative": on["e2e_latency_mean_s"]
            / off["e2e_latency_mean_s"]
            - 1,
            "e2e_p95_latency_relative": on["e2e_latency_p95_s"]
            / off["e2e_latency_p95_s"]
            - 1,
        },
        "gpu_hours_increment": state["gpu_hours_increment"],
        "gpu_hours_total": PRIOR_GPU_HOURS + state["gpu_hours_increment"],
    }
    common.atomic_json(ROOT / "result.json", result)
    state["status"] = "complete"
    state["result"] = result
    save_state(state)
    print(json.dumps(result, sort_keys=True))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("command", choices=("plan", "run"))
    args = parser.parse_args()
    if args.command == "plan":
        print(json.dumps(plan(), indent=2, sort_keys=True))
    else:
        run()


if __name__ == "__main__":
    main()