#!/usr/bin/env python3
"""Freeze Frontier's Qwen30 exact production-trace response surface."""

from __future__ import annotations

import argparse
import csv
import importlib.util
import json
import math
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Any


TARGET_PASS_RATE = 0.95
TPOT_SLOS_MS = (50.0, 100.0, 150.0, 180.0)
WINDOW_SECONDS = 600.0
BASE_RUNNER = (
    Path(__file__).resolve().parents[1]
    / "frontier-phase-factorial-v0/run_frontier_qwen30_prefill_surface.py"
)


def load_base():
    spec = importlib.util.spec_from_file_location("qwen30_prefill_surface_base", BASE_RUNNER)
    if spec is None or spec.loader is None:
        raise ImportError(BASE_RUNNER)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


BASE = load_base()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--frontier-source", type=Path, required=True)
    parser.add_argument("--replayserve-root", type=Path, required=True)
    parser.add_argument("--profile-root", type=Path, required=True)
    parser.add_argument("--python-deps", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument(
        "--trace",
        action="append",
        required=True,
        help="Frozen trace anchor as LABEL=PATH; repeat in increasing load order.",
    )
    parser.add_argument("--config", action="append")
    parser.add_argument(
        "--rate-contract",
        choices=("trace-window", "uniform-spacing"),
        default="trace-window",
    )
    parser.add_argument(
        "--prefix-caching",
        action=argparse.BooleanOptionalAction,
        default=True,
    )
    parser.add_argument(
        "--cc-backend", choices=("analytical", "vidur"), default="vidur"
    )
    parser.add_argument("--allreduce-csv", type=Path)
    parser.add_argument("--timeout-seconds", type=float, default=1800.0)
    parser.add_argument("--resume", action="store_true")
    return parser.parse_args()


def ttft_limit_ms(input_tokens: int) -> float:
    return 1000.0 + 1000.0 * input_tokens / 8000.0


def percentile(values: list[float], fraction: float) -> float | None:
    if not values:
        return None
    ordered = sorted(values)
    return ordered[math.ceil(fraction * len(ordered)) - 1]


def parse_trace(
    specification: str, *, rate_contract: str = "trace-window"
) -> dict[str, Any]:
    if "=" not in specification:
        raise ValueError(f"trace must be LABEL=PATH: {specification}")
    label, raw_path = specification.split("=", 1)
    if not label or "/" in label:
        raise ValueError(f"invalid trace label: {label!r}")
    path = Path(raw_path).resolve()
    if not path.is_file():
        raise FileNotFoundError(path)
    with path.open(newline="") as source:
        rows = list(csv.DictReader(source))
    if not rows:
        raise ValueError(f"empty trace: {path}")
    required = {
        "arrived_at",
        "num_prefill_tokens",
        "num_decode_tokens",
        "session_id",
        "block_hash_ids",
    }
    if not required.issubset(rows[0]):
        raise ValueError(f"trace columns missing: {required - set(rows[0])}")
    arrivals = [float(row["arrived_at"]) for row in rows]
    if any(not math.isfinite(value) or value < 0 for value in arrivals):
        raise ValueError(f"invalid arrival in {path}")
    if any(right < left for left, right in zip(arrivals, arrivals[1:])):
        raise ValueError(f"arrival order drift in {path}")
    if rate_contract == "trace-window":
        offered_request_rate = len(rows) / WINDOW_SECONDS
    elif rate_contract == "uniform-spacing":
        if len(rows) < 2 or arrivals[-1] <= arrivals[0]:
            raise ValueError("uniform-spacing traces require at least two arrivals")
        intervals = [right - left for left, right in zip(arrivals, arrivals[1:])]
        expected_interval = (arrivals[-1] - arrivals[0]) / (len(arrivals) - 1)
        if any(abs(value - expected_interval) > 1e-9 for value in intervals):
            raise ValueError(f"non-uniform fixed trace: {path}")
        offered_request_rate = 1.0 / expected_interval
    else:
        raise ValueError(f"unknown rate contract: {rate_contract}")
    shapes = [
        (int(row["num_prefill_tokens"]), int(row["num_decode_tokens"]))
        for row in rows
    ]
    if any(prefill <= 0 or decode <= 0 or prefill + decode > 40960 for prefill, decode in shapes):
        raise ValueError(f"out-of-contract shape in {path}")
    return {
        "label": label,
        "path": path,
        "sha256": BASE.sha256(path),
        "requests": len(rows),
        "offered_request_rate": offered_request_rate,
        "first_arrival_s": arrivals[0],
        "last_arrival_s": arrivals[-1],
        "shapes": shapes,
    }


def find_request_metrics(run_dir: Path) -> Path:
    matches = list((run_dir / "metrics").rglob("request_metrics.csv"))
    if len(matches) != 1:
        raise RuntimeError(f"expected one request_metrics.csv, got {matches}")
    return matches[0]


def trace_manifest_entry(trace: dict[str, Any]) -> dict[str, Any]:
    return {
        key: str(value) if isinstance(value, Path) else value
        for key, value in trace.items()
        if key != "shapes"
    }


def score(path: Path, expected_shapes: list[tuple[int, int]]) -> dict[str, Any]:
    with path.open(newline="") as source:
        rows = list(csv.DictReader(source))
    if len(rows) != len(expected_shapes):
        raise ValueError(f"request count mismatch: {len(rows)} != {len(expected_shapes)}")
    rows.sort(key=lambda row: int(row["Request Id"]))
    request_metrics = []
    for index, (row, expected) in enumerate(zip(rows, expected_shapes, strict=True)):
        if int(row["Request Id"]) != index:
            raise ValueError("Frontier request ID/order drift")
        shape = (
            int(float(row["request_num_prefill_tokens"])),
            int(float(row["request_num_decode_tokens"])),
        )
        if shape != expected:
            raise ValueError(f"request shape drift at {index}: {shape} != {expected}")
        ttft = float(row["ttft"])
        e2e = float(row["request_e2e_time"])
        tpot = (e2e - ttft) / (shape[1] - 1) if shape[1] > 1 else None
        values = [ttft, e2e] + ([] if tpot is None else [tpot])
        if not all(math.isfinite(value) and value >= 0 for value in values):
            raise ValueError(f"invalid latency at request {index}")
        request_metrics.append(
            {
                "request_id": index,
                "input_tokens": shape[0],
                "output_tokens": shape[1],
                "ttft_ms": ttft,
                "ttft_limit_ms": ttft_limit_ms(shape[0]),
                "tpot_ms": tpot,
                "e2e_ms": e2e,
            }
        )
    slos = {}
    for limit in TPOT_SLOS_MS:
        passed = sum(
            row["ttft_ms"] <= row["ttft_limit_ms"]
            and (row["tpot_ms"] is None or row["tpot_ms"] <= limit)
            for row in request_metrics
        )
        pass_rate = passed / len(request_metrics)
        slos[f"tpot_{int(limit)}ms"] = {
            "passed": passed,
            "pass_rate": pass_rate,
            "feasible": pass_rate >= TARGET_PASS_RATE,
        }
    ttfts = [float(row["ttft_ms"]) for row in request_metrics]
    tpots = [float(row["tpot_ms"]) for row in request_metrics if row["tpot_ms"] is not None]
    return {
        "ttft_p50_ms": percentile(ttfts, 0.50),
        "ttft_p95_ms": percentile(ttfts, 0.95),
        "tpot_p50_ms": percentile(tpots, 0.50),
        "tpot_p95_ms": percentile(tpots, 0.95),
        "slos": slos,
    }


def main() -> None:
    args = parse_args()
    for name in (
        "frontier_source",
        "replayserve_root",
        "profile_root",
        "python_deps",
        "output_root",
    ):
        setattr(args, name, getattr(args, name).resolve())
    if args.allreduce_csv is not None:
        args.allreduce_csv = args.allreduce_csv.resolve()
    traces = [
        parse_trace(specification, rate_contract=args.rate_contract)
        for specification in args.trace
    ]
    if len({trace["label"] for trace in traces}) != len(traces):
        raise ValueError("trace labels must be unique")
    if any(
        right["offered_request_rate"] <= left["offered_request_rate"]
        for left, right in zip(traces, traces[1:])
    ):
        raise ValueError("trace anchors must be supplied in increasing load order")

    selected = list(BASE.GRID)
    if args.config:
        wanted = set(args.config)
        selected = [config for config in BASE.GRID if config.name in wanted]
        if {config.name for config in selected} != wanted:
            raise ValueError(f"unknown configs: {wanted - {config.name for config in selected}}")
    paths = BASE.profile_paths(args.profile_root)
    coverage = BASE.validate_profile(paths)
    builder = BASE.load_module(
        "qwen30_exact_trace_frontier_builder",
        args.replayserve_root / "tools/run_frontier_sweep.py",
    )
    frontier_head = subprocess.run(
        ["git", "-C", str(args.frontier_source), "rev-parse", "HEAD"],
        check=True,
        text=True,
        stdout=subprocess.PIPE,
    ).stdout.strip()
    environment = os.environ.copy()
    pythonpath = [str(args.python_deps), str(args.frontier_source)]
    if environment.get("PYTHONPATH"):
        pythonpath.append(environment["PYTHONPATH"])
    environment.update(
        {
            "PYTHONPATH": ":".join(pythonpath),
            "CUDA_VISIBLE_DEVICES": "",
            "NVIDIA_VISIBLE_DEVICES": "void",
            "WANDB_DISABLED": "true",
            "VIDUR_DISABLE_WANDB": "1",
            "FRONTIER_LOG_LEVEL": "WARNING",
            "PYTHONDONTWRITEBYTECODE": "1",
        }
    )

    config_results = []
    for config in selected:
        loads = []
        config_knobs = BASE.knobs(config, paths, args.output_root / "cache")
        config_knobs["enable_prefix_caching"] = args.prefix_caching
        config_knobs["prediction_max_tokens_per_request"] = 40960
        for trace in traces:
            run_dir = args.output_root / "runs" / config.name / trace["label"]
            result_path = run_dir / "result.json"
            if args.resume and result_path.is_file():
                loads.append(json.loads(result_path.read_text()))
                continue
            run_dir.mkdir(parents=True, exist_ok=True)
            command = builder.build_frontier_command(
                python_bin="/usr/bin/python3",
                trace_file=trace["path"],
                metrics_root=run_dir / "metrics",
                run_id=f"qwen30_trace_{config.name}_{trace['label']}",
                knobs=config_knobs,
            )
            command = BASE.configure_cc_command(
                command,
                backend=args.cc_backend,
                allreduce_csv=args.allreduce_csv,
                cache=args.output_root / "cc-cache",
            )
            BASE.write_json(run_dir / "command.json", command)
            started = time.time()
            with (run_dir / "stdout.log").open("w") as stdout, (
                run_dir / "stderr.log"
            ).open("w") as stderr:
                completed = subprocess.run(
                    command,
                    cwd=args.frontier_source,
                    env=environment,
                    stdout=stdout,
                    stderr=stderr,
                    timeout=args.timeout_seconds,
                    check=False,
                )
            if completed.returncode != 0:
                raise RuntimeError(
                    f"Frontier failed for {config.name}/{trace['label']}: {completed.returncode}"
                )
            metrics = find_request_metrics(run_dir)
            result = {
                "status": "completed",
                "config": {"tp": config.tp, "mns": config.mns, "name": config.name},
                "trace_label": trace["label"],
                "offered_request_rate": trace["offered_request_rate"],
                "offered_request_rate_per_gpu": trace["offered_request_rate"] / config.tp,
                "request_count": trace["requests"],
                "elapsed_seconds": time.time() - started,
                "trace_sha256": trace["sha256"],
                "request_metrics_sha256": BASE.sha256(metrics),
                "score": score(metrics, trace["shapes"]),
            }
            BASE.write_json(result_path, result)
            loads.append(result)
            print(
                json.dumps(
                    {
                        "config": config.name,
                        "trace": trace["label"],
                        "rate": trace["offered_request_rate"],
                        "primary": result["score"]["slos"]["tpot_150ms"],
                    },
                    sort_keys=True,
                ),
                flush=True,
            )
        config_results.append(
            {
                "config": {"tp": config.tp, "mns": config.mns, "name": config.name},
                "loads": loads,
            }
        )

    rankings = {}
    for slo in (f"tpot_{int(value)}ms" for value in TPOT_SLOS_MS):
        records = []
        for item in config_results:
            feasible = [
                load["offered_request_rate"]
                for load in item["loads"]
                if load["score"]["slos"][slo]["feasible"]
            ]
            capacity = max(feasible, default=None)
            records.append(
                {
                    "config": item["config"],
                    "maximum_tested_feasible_request_rate": capacity,
                    "maximum_tested_feasible_request_rate_per_gpu": (
                        capacity / item["config"]["tp"] if capacity is not None else None
                    ),
                    "lower_censored": capacity is None,
                    "upper_censored": capacity == traces[-1]["offered_request_rate"],
                }
            )
        records.sort(
            key=lambda row: (
                -(
                    row["maximum_tested_feasible_request_rate_per_gpu"]
                    if row["maximum_tested_feasible_request_rate_per_gpu"] is not None
                    else -1
                ),
                row["config"]["name"],
            )
        )
        rankings[slo] = records

    manifest = {
        "schema": "frontier-qwen30-exact-trace-surface-v1",
        "status": "frozen_before_real" if selected == list(BASE.GRID) else "partial_not_decision_bearing",
        "contract": {
            "window_seconds": (
                WINDOW_SECONDS if args.rate_contract == "trace-window" else None
            ),
            "rate_contract": args.rate_contract,
            "prefix_caching": args.prefix_caching,
            "arrival": "original_trace_timestamp_and_order",
            "input_output": "exact_source_values",
            "ttft_slo": "1000ms + 1000ms * input_tokens / 8000",
            "tpot_slos_ms": TPOT_SLOS_MS,
            "primary_tpot_slo_ms": 150.0,
            "target_pass_rate": TARGET_PASS_RATE,
        },
        "frontier": {
            "source": str(args.frontier_source),
            "git_head": frontier_head,
            "git_status_short": subprocess.run(
                ["git", "-C", str(args.frontier_source), "status", "--short"],
                check=True,
                text=True,
                stdout=subprocess.PIPE,
            ).stdout,
        },
        "profiles": {
            "root": str(args.profile_root),
            "coverage": coverage,
            "sha256": {name: BASE.sha256(path) for name, path in paths.items()},
        },
        "collective": {
            "backend": args.cc_backend,
            "allreduce_csv": str(args.allreduce_csv) if args.allreduce_csv else None,
            "allreduce_csv_sha256": BASE.sha256(args.allreduce_csv) if args.allreduce_csv else None,
        },
        "traces": [trace_manifest_entry(trace) for trace in traces],
        "config_results": config_results,
        "rankings": rankings,
    }
    BASE.write_json(args.output_root / "frontier_trace_surface_frozen.json", manifest)
    print(args.output_root / "frontier_trace_surface_frozen.json")


if __name__ == "__main__":
    main()