#!/usr/bin/env python3
"""Freeze the Qwen30 fixed-shape prefill-only Frontier surface."""

from __future__ import annotations

import argparse
import csv
import hashlib
import importlib.util
import json
import math
import os
import subprocess
import sys
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any


MODEL = "qwen3-a3b-30b-moe"
RATES = (4.0, 8.0, 16.0, 32.0, 64.0)
TTFT_SLO_MS = 1256.0
TARGET_PASS_RATE = 0.95
NUM_BLOCKS = {1: 20080, 2: 76537, 4: 191727}


@dataclass(frozen=True)
class Config:
    tp: int
    mns: int

    @property
    def name(self) -> str:
        return f"tp{self.tp}_mns{self.mns}"


GRID = tuple(Config(tp, mns) for tp in (1, 2, 4) for mns in (8, 16, 32, 64))


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("--requests", type=int, default=64)
    parser.add_argument("--rate", type=float, action="append")
    parser.add_argument("--config", action="append")
    parser.add_argument(
        "--cc-backend", choices=("analytical", "vidur"), default="analytical"
    )
    parser.add_argument("--allreduce-csv", type=Path)
    parser.add_argument("--timeout-seconds", type=float, default=900.0)
    parser.add_argument("--resume", action="store_true")
    return parser.parse_args()


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    os.replace(temporary, path)


def load_module(name: str, path: Path):
    spec = importlib.util.spec_from_file_location(name, path)
    if spec is None or spec.loader is None:
        raise ImportError(path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


def write_trace(path: Path, *, requests: int, rate: float) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fields = [
        "arrived_at",
        "num_prefill_tokens",
        "num_decode_tokens",
        "session_id",
        "block_hash_ids",
    ]
    with path.open("w", newline="") as output:
        writer = csv.DictWriter(output, fieldnames=fields)
        writer.writeheader()
        for request_id in range(requests):
            writer.writerow(
                {
                    "arrived_at": f"{request_id / rate:.12f}",
                    "num_prefill_tokens": 2048,
                    "num_decode_tokens": 1,
                    "session_id": request_id,
                    "block_hash_ids": "|".join(
                        str(request_id * 128 + block + 1) for block in range(128)
                    ),
                }
            )


def profile_paths(root: Path) -> dict[str, Path]:
    paths = {
        "linear": root / "linear_op.csv",
        "attention": root / "attention.csv",
        "moe": root / "moe.csv",
        "manifest": root / "manifest.json",
    }
    missing = [str(path) for path in paths.values() if not path.is_file()]
    if missing:
        raise FileNotFoundError(missing)
    return paths


def validate_profile(paths: dict[str, Path]) -> dict[str, Any]:
    manifest = json.loads(paths["manifest"].read_text())
    expected = manifest["outputs"]
    for filename in ("linear_op.csv", "attention.csv", "moe.csv"):
        actual = sha256(paths[{"linear_op.csv": "linear", "attention.csv": "attention", "moe.csv": "moe"}[filename]])
        if actual != expected[filename]:
            raise ValueError(f"profile hash mismatch for {filename}")
    with paths["attention"].open(newline="") as source:
        rows = list(csv.DictReader(source))
    coverage = {}
    for tp in (1, 2, 4):
        exact = [
            row
            for row in rows
            if int(row["num_tensor_parallel_workers"]) == tp
            and row["is_prefill"].lower() == "true"
            and row.get("is_true_mixed_batch", "").lower() != "true"
            and int(float(row["batch_size"])) == 1
            and int(float(row["total_tokens"])) == 2048
        ]
        if len(exact) != 1:
            raise ValueError(f"expected one exact TP{tp} 2048-token prefill row, got {len(exact)}")
        coverage[str(tp)] = {"exact_prefill_2048_rows": len(exact), "profile_batch_size": int(exact[0]["batch_size"])}
    return {"manifest": manifest, "attention": coverage}


def knobs(config: Config, paths: dict[str, Path], cache: Path) -> dict[str, Any]:
    return {
        "simulation_mode": "online",
        "sys_arch": "co-location",
        "num_replicas": 1,
        "cluster_scheduler": "sticky_round_robin",
        "model_name": MODEL,
        "device": "h20",
        "network_device": "h20_dgx",
        "attn_tensor_parallel_size": config.tp,
        "attn_data_parallel_size": 1,
        "moe_tensor_parallel_size": config.tp,
        "moe_expert_parallel_size": 1,
        "num_pipeline_stages": 1,
        "replica_scheduler": "vllm_v1",
        "decode_cuda_graph_mode": "none",
        "batch_size_cap": config.mns,
        "max_tokens_in_batch": 8192,
        "long_prefill_token_threshold": 0,
        "block_size": 16,
        "num_blocks_mode": "explicit",
        "num_blocks": NUM_BLOCKS[config.tp],
        "gpu_memory_utilization": 0.92,
        "non_kv_cache_overhead_bytes": 0,
        "trace_max_tokens": 40960,
        "cache_dir": str(cache),
        "enable_dummy_mode": False,
        "linear_op_input_file": str(paths["linear"]),
        "atten_input_file": str(paths["attention"]),
        "moe_input_file": str(paths["moe"]),
        "prediction_max_prefill_chunk_size": 18000,
        "prediction_max_batch_size": 128,
        "prediction_max_tokens_per_request": 32768,
        "no_cache": False,
        "skip_cpu_overhead_modeling": True,
        "enable_prefix_caching": False,
        "enable_chunked_prefill": True,
    }


def configure_cc_command(
    command: list[str], *, backend: str, allreduce_csv: Path | None, cache: Path
) -> list[str]:
    configured = list(command)
    option = "--cc_backend_config_type"
    try:
        index = configured.index(option)
    except ValueError as error:
        raise ValueError(f"Frontier command is missing {option}") from error
    configured[index + 1] = backend
    if backend == "analytical":
        if allreduce_csv is not None:
            raise ValueError("--allreduce-csv requires --cc-backend vidur")
        return configured
    if allreduce_csv is None:
        raise ValueError("--cc-backend vidur requires --allreduce-csv")
    configured.extend(
        [
            "--vidur_cc_backend_config_all_reduce_input_file",
            str(allreduce_csv),
            "--vidur_cc_backend_config_cache_dir",
            str(cache),
            "--vidur_cc_backend_config_k_fold_cv_splits",
            "6",
            "--vidur_cc_backend_config_num_training_job_threads",
            "1",
        ]
    )
    return configured


def find_metrics(run_dir: Path) -> tuple[Path, Path]:
    systems = list((run_dir / "metrics").rglob("system_metrics.json"))
    requests = list((run_dir / "metrics").rglob("request_metrics.csv"))
    if len(systems) != 1 or len(requests) != 1:
        raise RuntimeError(f"metric pair mismatch: {len(systems)}/{len(requests)}")
    return systems[0], requests[0]


def score(system_path: Path, request_path: Path, expected_requests: int) -> dict[str, Any]:
    system = json.loads(system_path.read_text())
    metadata = system["simulation_metadata"]
    if int(metadata["completed_requests"]) != expected_requests:
        raise ValueError("Frontier completion count mismatch")
    with request_path.open(newline="") as source:
        rows = list(csv.DictReader(source))
    if len(rows) != expected_requests:
        raise ValueError("request metric count mismatch")
    values = []
    passed = 0
    for row in rows:
        if int(float(row["request_num_prefill_tokens"])) != 2048:
            raise ValueError("prefill shape drift")
        if int(float(row["request_num_decode_tokens"])) != 1:
            raise ValueError("decode shape drift")
        ttft = float(row["ttft"])
        if not math.isfinite(ttft) or ttft < 0:
            raise ValueError("invalid TTFT")
        values.append(ttft)
        passed += int(ttft <= TTFT_SLO_MS)
    ordered = sorted(values)
    pass_rate = passed / expected_requests
    return {
        "ttft_p50_ms": ordered[math.ceil(0.50 * len(ordered)) - 1],
        "ttft_p95_ms": ordered[math.ceil(0.95 * len(ordered)) - 1],
        "ttft_max_ms": max(ordered),
        "passed": passed,
        "pass_rate": pass_rate,
        "feasible": pass_rate >= TARGET_PASS_RATE,
        "throughput_requests_per_second": float(system["throughput_metrics"]["requests_per_second"]),
    }


def main() -> None:
    args = parse_args()
    args.frontier_source = args.frontier_source.resolve()
    args.replayserve_root = args.replayserve_root.resolve()
    args.profile_root = args.profile_root.resolve()
    args.python_deps = args.python_deps.resolve()
    args.output_root = args.output_root.resolve()
    if args.allreduce_csv is not None:
        args.allreduce_csv = args.allreduce_csv.resolve()
        if not args.allreduce_csv.is_file():
            raise FileNotFoundError(args.allreduce_csv)
    rates = tuple(args.rate or RATES)
    selected = list(GRID)
    if args.config:
        wanted = set(args.config)
        selected = [config for config in 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 = profile_paths(args.profile_root)
    coverage = validate_profile(paths)
    builder = load_module(
        "qwen30_prefill_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()
    traces = {}
    for rate in rates:
        trace = args.output_root / "traces" / f"r{rate:g}.csv"
        write_trace(trace, requests=args.requests, rate=rate)
        traces[rate] = trace

    config_results = []
    for config in selected:
        loads = []
        config_knobs = knobs(config, paths, args.output_root / "cache")
        for rate in rates:
            run_dir = args.output_root / "runs" / config.name / f"r{rate:g}"
            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=traces[rate],
                metrics_root=run_dir / "metrics",
                run_id=f"qwen30_prefill_{config.name}_r{rate:g}",
                knobs=config_knobs,
            )
            command = configure_cc_command(
                command,
                backend=args.cc_backend,
                allreduce_csv=args.allreduce_csv,
                cache=args.output_root / "cc-cache",
            )
            write_json(run_dir / "command.json", command)
            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",
                }
            )
            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} rate={rate}: {completed.returncode}"
                )
            system_path, request_path = find_metrics(run_dir)
            result = {
                "status": "completed",
                "config": asdict(config) | {"name": config.name},
                "offered_request_rate": rate,
                "offered_request_rate_per_gpu": rate / config.tp,
                "request_count": args.requests,
                "elapsed_seconds": time.time() - started,
                "trace_sha256": sha256(traces[rate]),
                "request_metrics_sha256": sha256(request_path),
                "score": score(system_path, request_path, args.requests),
            }
            write_json(result_path, result)
            loads.append(result)
            print(
                json.dumps(
                    {
                        "config": config.name,
                        "rate": rate,
                        "pass_rate": result["score"]["pass_rate"],
                        "feasible": result["score"]["feasible"],
                    },
                    sort_keys=True,
                ),
                flush=True,
            )
        config_results.append({"config": asdict(config) | {"name": config.name}, "loads": loads})

    capacities = []
    for item in config_results:
        feasible = [
            load["offered_request_rate"]
            for load in item["loads"]
            if load["score"]["feasible"]
        ]
        capacity = max(feasible) if feasible else None
        capacities.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 == max(rates) if capacity is not None else False,
            }
        )
    capacities.sort(
        key=lambda row: (
            -(row["maximum_tested_feasible_request_rate_per_gpu"] or -1),
            row["config"]["name"],
        )
    )
    full = selected == list(GRID) and rates == RATES and args.requests == 64
    manifest = {
        "schema": "frontier-qwen30-prefill-surface-v1",
        "status": "frozen_before_real" if full else "partial_not_decision_bearing",
        "contract": {
            "rates": rates,
            "requests_per_anchor": args.requests,
            "input_tokens": 2048,
            "output_tokens": 1,
            "ttft_slo_ms": TTFT_SLO_MS,
            "target_pass_rate": TARGET_PASS_RATE,
            "prefix_caching": False,
            "arrival": "open_loop_uniform",
        },
        "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: sha256(path) for name, path in paths.items()},
        },
        "collective": {
            "backend": args.cc_backend,
            "allreduce_csv": (
                str(args.allreduce_csv) if args.allreduce_csv is not None else None
            ),
            "allreduce_csv_sha256": (
                sha256(args.allreduce_csv) if args.allreduce_csv is not None else None
            ),
        },
        "config_results": config_results,
        "capacity": capacities,
    }
    write_json(args.output_root / "frontier_surface_frozen.json", manifest)
    print(args.output_root / "frontier_surface_frozen.json")


if __name__ == "__main__":
    main()
