#!/usr/bin/env python3
"""Freeze the full fixed-shape T0 Frontier response surface."""

from __future__ import annotations

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


MODEL = "Qwen3-235B-A22B-FP8"
FRONTIER_DECLARED_UPSTREAM_COMMIT = "d9cfeb6d8791fbf2f295dd9744c56a666171776e"
FRONTIER_PYTHON_CONFIG_TREE_SHA256 = "172fc7ae19c40e67030208ff488d0d5d90764888ed76a7a543be6037ba62dc11"
RATES = (0.10, 0.20, 0.40, 0.80, 1.20, 1.60, 2.40, 3.20)
TPOT_SLOS_MS = (40.0, 120.0, 150.0, 180.0)
TTFT_SLO_MS = 1256.0
TARGET_PASS_RATE = 0.95


@dataclass(frozen=True)
class Config:
    tp: int
    mns: int
    mbt: int
    moe_tp: int
    moe_ep: int
    num_gpu_blocks: int

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


GRID = tuple(
    Config(
        tp=tp,
        mns=mns,
        mbt=mbt,
        moe_tp=4 if tp == 4 else 1,
        moe_ep=1 if tp == 4 else 8,
        num_gpu_blocks=26101 if tp == 4 else 62351,
    )
    for tp in (4, 8)
    for mns in (64, 128)
    for mbt in (8192, 16384)
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--frontier-source", type=Path, required=True)
    parser.add_argument("--profile-root", type=Path, required=True)
    parser.add_argument("--python", 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("--subprocess-timeout-seconds", type=int, default=1800)
    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 rate_key(rate: float) -> str:
    return f"r{rate:.2f}".replace(".", "p")


def write_trace(path: Path, *, request_count: int, rate: float) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as output:
        writer = csv.DictWriter(
            output,
            fieldnames=[
                "arrived_at",
                "num_prefill_tokens",
                "num_decode_tokens",
                "slo_ttft_ms",
            ],
        )
        writer.writeheader()
        for request_id in range(request_count):
            writer.writerow(
                {
                    "arrived_at": f"{request_id / rate:.12f}",
                    "num_prefill_tokens": 2048,
                    "num_decode_tokens": 128,
                    "slo_ttft_ms": TTFT_SLO_MS,
                }
            )


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


def validate_attention_coverage(path: Path) -> dict[str, Any]:
    with path.open(newline="") as source:
        rows = list(csv.DictReader(source))
    standard_decode = [
        row
        for row in rows
        if row["is_prefill"].lower() == "false"
        and row.get("is_true_mixed_batch", "").lower() != "true"
    ]
    true_mixed = [
        row for row in rows if row.get("is_true_mixed_batch", "").lower() == "true"
    ]
    standard_by_tp = {
        str(tp): sum(int(row["num_tensor_parallel_workers"]) == tp for row in standard_decode)
        for tp in (4, 8)
    }
    mixed_by_tp = {
        str(tp): sum(int(row["num_tensor_parallel_workers"]) == tp for row in true_mixed)
        for tp in (4, 8)
    }
    if standard_by_tp != {"4": 81, "8": 81}:
        raise ValueError(f"standard decode coverage mismatch: {standard_by_tp}")
    if mixed_by_tp != {"4": 108, "8": 108}:
        raise ValueError(f"true-mixed coverage mismatch: {mixed_by_tp}")
    return {
        "rows": len(rows),
        "standard_decode_by_tp": standard_by_tp,
        "true_mixed_by_tp": mixed_by_tp,
    }


def build_command(
    args: argparse.Namespace,
    paths: dict[str, Path],
    config: Config,
    trace: Path,
    run_dir: Path,
) -> list[str]:
    cache = args.output_root / "cache" / f"tp{config.tp}"
    return [
        str(args.python), "-m", "frontier.main",
        "--simulation_mode", "offline",
        "--offline_use_generated_request_arrivals",
        "--sys_arch", "co-location",
        "--cluster_config_num_replicas", "1",
        "--replica_config_model_name", MODEL,
        "--replica_config_attn_tensor_parallel_size", str(config.tp),
        "--replica_config_attn_data_parallel_size", "1",
        "--replica_config_moe_tensor_parallel_size", str(config.moe_tp),
        "--replica_config_moe_expert_parallel_size", str(config.moe_ep),
        "--replica_config_total_expert_num", "128",
        "--replica_config_router_topk", "8",
        "--replica_config_moe_routing_mode", "simulation",
        "--replica_config_moe_routing_seed", "42",
        "--replica_config_num_pipeline_stages", "1",
        "--replica_config_device", "h20",
        "--replica_config_network_device", "h20_dgx",
        "--cc_backend_config_type", "vidur",
        "--vidur_cc_backend_config_profiling_data_dir", str(args.profile_root),
        "--vidur_cc_backend_config_cache_dir", str(cache / "collectives"),
        "--vidur_cc_backend_config_all_reduce_input_file", str(paths["all_reduce"]),
        "--replica_scheduler_config_type", "vllm_v1",
        "--decode_cuda_graph_mode", "none",
        "--vllm_v1_scheduler_config_batch_size_cap", str(config.mns),
        "--vllm_v1_scheduler_config_block_size", "16",
        "--vllm_v1_scheduler_config_num_blocks", str(config.num_gpu_blocks),
        "--vllm_v1_scheduler_config_num_blocks_mode", "explicit",
        "--vllm_v1_scheduler_config_max_tokens_in_batch", str(config.mbt),
        "--vllm_v1_scheduler_config_enable_chunked_prefill",
        "--no-vllm_v1_scheduler_config_enable_prefix_caching",
        "--request_generator_config_type", "trace_replay",
        "--trace_request_generator_config_trace_file", str(trace),
        "--trace_request_generator_config_time_scale_factor", "1",
        "--trace_request_generator_config_prefill_scale_factor", "1",
        "--trace_request_generator_config_decode_scale_factor", "1",
        "--trace_request_generator_config_max_tokens", "40960",
        "--no-random_forrest_execution_time_predictor_config_enable_dummy_mode",
        "--random_forrest_execution_time_predictor_config_linear_op_input_file", str(paths["linear"]),
        "--random_forrest_execution_time_predictor_config_atten_input_file", str(paths["attention"]),
        "--random_forrest_execution_time_predictor_config_moe_input_file", str(paths["moe"]),
        "--random_forrest_execution_time_predictor_config_all_reduce_input_file", str(paths["all_reduce"]),
        "--random_forrest_execution_time_predictor_config_prediction_max_prefill_chunk_size", "16384",
        "--random_forrest_execution_time_predictor_config_prediction_max_tokens_per_request", "40960",
        "--random_forrest_execution_time_predictor_config_prediction_max_batch_size", "128",
        "--random_forrest_execution_time_predictor_config_skip_cpu_overhead_modeling",
        "--metrics_config_cache_dir", str(cache / "execution"),
        "--metrics_config_output_dir", str(run_dir / "metrics"),
        "--metrics_config_run_id", f"{config.name}_{run_dir.name}",
        "--metrics_config_write_metrics",
        "--metrics_config_store_request_metrics",
        "--no-metrics_config_store_plots",
        "--no-metrics_config_enable_chrome_trace",
        "--no-metrics_config_write_json_trace",
    ]


def find_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 score(metrics: Path, expected_requests: int) -> dict[str, Any]:
    with metrics.open(newline="") as source:
        rows = list(csv.DictReader(source))
    if len(rows) != expected_requests:
        raise ValueError(f"request count mismatch: {len(rows)} != {expected_requests}")
    requests = []
    for row in rows:
        prompt = int(float(row["request_num_prefill_tokens"]))
        completion = int(float(row["request_num_decode_tokens"]))
        ttft = float(row["ttft"])
        e2e = float(row["request_e2e_time"])
        tpot = (e2e - ttft) / (completion - 1)
        if prompt != 2048 or completion != 128:
            raise ValueError("request shape drift")
        if not all(math.isfinite(value) and value >= 0 for value in (ttft, e2e, tpot)):
            raise ValueError("non-finite or negative latency")
        requests.append({"request_id": int(row["Request Id"]), "ttft_ms": ttft, "tpot_ms": tpot, "e2e_ms": e2e})
    slos = {}
    for limit in TPOT_SLOS_MS:
        passed = sum(row["ttft_ms"] <= TTFT_SLO_MS and row["tpot_ms"] <= limit for row in requests)
        pass_rate = passed / len(requests)
        slos[f"tpot_{int(limit)}ms"] = {
            "passed": passed,
            "pass_rate": pass_rate,
            "feasible": pass_rate >= TARGET_PASS_RATE,
        }
    return {"requests": requests, "slos": slos}


def main() -> None:
    args = parse_args()
    args.frontier_source = args.frontier_source.resolve()
    args.profile_root = args.profile_root.resolve()
    # A venv interpreter is commonly a symlink to the system executable.
    # Keep the venv path so Python discovers that environment's site-packages.
    args.python = args.python.absolute()
    args.output_root = args.output_root.resolve()
    if args.requests < 2:
        raise ValueError("requests must be at least two")
    rates = tuple(args.rate or RATES)
    if any(rate <= 0 for rate in rates) or len(set(rates)) != len(rates):
        raise ValueError("rates must be positive and unique")
    selected = list(GRID)
    if args.config:
        names = set(args.config)
        selected = [config for config in GRID if config.name in names]
        if {config.name for config in selected} != names:
            raise ValueError(f"unknown configs: {names - {config.name for config in selected}}")

    paths = profile_paths(args.profile_root)
    coverage = validate_attention_coverage(paths["attention"])
    args.output_root.mkdir(parents=True, exist_ok=True)
    traces = {}
    for rate in rates:
        path = args.output_root / "traces" / f"{rate_key(rate)}.csv"
        write_trace(path, request_count=args.requests, rate=rate)
        traces[rate] = path

    config_results = []
    environment = dict(os.environ)
    environment.update(
        {
            "PYTHONPATH": str(args.frontier_source),
            "WANDB_DISABLED": "true",
            "VIDUR_DISABLE_WANDB": "1",
            # The best-effort source emits a per-layer OP-TRACE at INFO.  It is
            # diagnostic only and can produce hundreds of MiB per T0 cell.
            "FRONTIER_LOG_LEVEL": "WARNING",
        }
    )
    for config in selected:
        loads = []
        for rate in rates:
            run_dir = args.output_root / "runs" / config.name / rate_key(rate)
            result_path = run_dir / "result.json"
            if result_path.is_file():
                result = json.loads(result_path.read_text())
                if result.get("status") == "completed":
                    loads.append(result)
                    continue
            run_dir.mkdir(parents=True, exist_ok=True)
            command = build_command(args, paths, config, traces[rate], run_dir)
            write_json(run_dir / "command.json", command)
            started = time.time()
            with (run_dir / "stdout.log").open("w") as output:
                completed = subprocess.run(
                    command,
                    cwd=args.frontier_source,
                    env=environment,
                    stdout=output,
                    stderr=subprocess.STDOUT,
                    timeout=args.subprocess_timeout_seconds,
                    check=False,
                    text=True,
                )
            if completed.returncode != 0:
                raise RuntimeError(f"Frontier failed: {config.name} rate={rate}, rc={completed.returncode}")
            metrics = find_metrics(run_dir)
            result = {
                "status": "completed",
                "config": asdict(config) | {"name": config.name},
                "offered_request_rate": rate,
                "request_rate_per_gpu": rate / config.tp,
                "elapsed_seconds": time.time() - started,
                "trace_sha256": sha256(traces[rate]),
                "request_metrics_sha256": sha256(metrics),
                **score(metrics, args.requests),
            }
            write_json(result_path, result)
            loads.append(result)
            print(json.dumps({"config": config.name, "rate": rate, "elapsed_seconds": result["elapsed_seconds"], "slos": result["slos"]}, sort_keys=True), flush=True)
        config_results.append({"config": asdict(config) | {"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["slos"][slo]["feasible"]]
            capacity = max(feasible) if feasible else 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 == max(rates) if capacity is not None else False,
            })
        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

    is_complete_preregistered_surface = (
        selected == list(GRID) and rates == RATES and args.requests == 64
    )
    manifest = {
        "schema": "frontier-qwen235b-t0-surface-v1",
        "status": (
            "frozen_before_real_surface"
            if is_complete_preregistered_surface
            else "partial_surface_not_decision_bearing"
        ),
        "contract": {"requests_per_anchor": args.requests, "rates": rates, "input_tokens": 2048, "output_tokens": 128, "ttft_slo_ms": TTFT_SLO_MS, "tpot_slos_ms": TPOT_SLOS_MS, "target_pass_rate": TARGET_PASS_RATE, "prefix_caching": False},
        "frontier": {
            "source": str(args.frontier_source),
            "declared_upstream_commit": FRONTIER_DECLARED_UPSTREAM_COMMIT,
            "python_and_config_tree_sha256": FRONTIER_PYTHON_CONFIG_TREE_SHA256,
            "fingerprint_source": "prefill-grid-v3 frozen run manifest for the same immutable source snapshot",
        },
        "runner": {"path": str(Path(__file__).resolve()), "sha256": sha256(Path(__file__).resolve())},
        "profiles": {"root": str(args.profile_root), "coverage": coverage, "files_sha256": {name: sha256(path) for name, path in paths.items()}},
        "config_results": config_results,
        "rankings": rankings,
    }
    write_json(args.output_root / "frontier_surface_frozen.json", manifest)
    print(args.output_root / "frontier_surface_frozen.json")


if __name__ == "__main__":
    main()
