#!/usr/bin/env python3
"""Run Frontier on the four-cell Qwen235 vLLM 0.20 latency surface."""

from __future__ import annotations

import argparse
import importlib.util
import json
import os
import subprocess
import sys
import time
from dataclasses import asdict, dataclass
from pathlib import Path


HERE = Path(__file__).resolve().parent
Q30_RUNNER = HERE / "run_frontier_qwen30_exact_trace_surface.py"


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


Q30 = load_q30()
BASE = Q30.BASE
MODEL = "Qwen3-235B-A22B"


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

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


GRID = tuple(
    Config(tp, mns, 4 if tp == 4 else 1, 1 if tp == 4 else 8)
    for tp in (4, 8)
    for mns in (64, 128)
)


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("--runtime-contract", type=Path, required=True)
    parser.add_argument("--trace-tp", action="append", required=True, help="TP=PATH")
    parser.add_argument("--config", action="append")
    parser.add_argument("--prefix-caching", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--allreduce-csv", type=Path, required=True)
    parser.add_argument("--timeout-seconds", type=float, default=7200)
    parser.add_argument("--resume", action="store_true")
    return parser.parse_args()


def profile_paths(root: Path) -> dict[str, Path]:
    paths = {
        "linear": root / "linear_op.csv",
        "attention": root / "attention.csv",
        "moe": root / "moe.csv",
        "linear_kernel": root / "linear_op_kernel_only.csv",
        "attention_kernel": root / "attention_kernel_only.csv",
        "moe_kernel": root / "moe_kernel_only.csv",
        "manifest": root / "manifest.json",
    }
    missing = [str(path) for path in paths.values() if not path.is_file()]
    if missing:
        raise FileNotFoundError(missing)
    manifest = json.loads(paths["manifest"].read_text())
    for name, path in paths.items():
        if name == "manifest":
            continue
        expected = manifest.get("outputs", {}).get(path.name)
        if expected != BASE.sha256(path):
            raise ValueError(f"profile hash mismatch: {path}")
    return paths


def knobs(config: Config, paths: dict[str, Path], contract: dict, cache: Path, prefix: bool):
    resolved = contract[config.name]
    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.moe_tp,
        "moe_expert_parallel_size": config.moe_ep,
        "num_pipeline_stages": 1,
        "replica_scheduler": "vllm_v1",
        "decode_cuda_graph_mode": "piecewise",
        "batch_size_cap": config.mns,
        "max_tokens_in_batch": 8192,
        "long_prefill_token_threshold": 0,
        "enable_chunked_prefill": True,
        "block_size": 16,
        "num_blocks_mode": "explicit",
        "num_blocks": int(resolved["num_gpu_blocks"]),
        "gpu_memory_utilization": 0.80,
        "non_kv_cache_overhead_bytes": 0,
        "trace_max_tokens": 40960,
        "cache_dir": str(cache / config.name),
        "enable_prefix_caching": prefix,
        "enable_dummy_mode": False,
        "linear_op_input_file": str(paths["linear"]),
        "atten_input_file": str(paths["attention"]),
        "moe_input_file": str(paths["moe"]),
        "linear_op_kernel_only_input_file": str(paths["linear_kernel"]),
        "atten_kernel_only_input_file": str(paths["attention_kernel"]),
        "moe_kernel_only_input_file": str(paths["moe_kernel"]),
        "prediction_max_prefill_chunk_size": 8192,
        "prediction_max_tokens_per_request": 40960,
        "prediction_max_batch_size": max(int(v) for v in resolved["capture_sizes"]),
        "no_cache": True,
    }


def main() -> None:
    args = parse_args()
    for name in ("frontier_source", "replayserve_root", "profile_root", "python_deps", "output_root", "runtime_contract", "allreduce_csv"):
        setattr(args, name, getattr(args, name).resolve())
    paths = profile_paths(args.profile_root)
    contract = json.loads(args.runtime_contract.read_text())["configs"]
    trace_by_tp = {}
    for specification in args.trace_tp:
        raw_tp, separator, path = specification.partition("=")
        if not separator:
            raise ValueError(f"trace-tp must be TP=PATH: {specification}")
        tp = int(raw_tp)
        trace_by_tp[tp] = Q30.parse_trace(
            f"eval={path}", rate_contract="trace-window", prefix_caching=args.prefix_caching
        )
    if set(trace_by_tp) != {4, 8}:
        raise ValueError("trace-tp must provide exactly TP4 and TP8")
    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}}")
    allreduce_profile = Q30.validate_allreduce_profile(
        args.allreduce_csv, {config.tp for config in selected if config.tp > 1}
    )
    allreduce_csv_sha256 = BASE.sha256(args.allreduce_csv)

    builder = BASE.load_module(
        "qwen235_frontier_builder", args.replayserve_root / "tools/run_frontier_sweep.py"
    )
    environment = os.environ.copy()
    environment.update(
        {
            "PYTHONPATH": ":".join([str(args.python_deps), str(args.frontier_source)]),
            "CUDA_VISIBLE_DEVICES": "",
            "NVIDIA_VISIBLE_DEVICES": "void",
            "WANDB_DISABLED": "true",
            "VIDUR_DISABLE_WANDB": "1",
            "FRONTIER_LOG_LEVEL": "WARNING",
            "PYTHONDONTWRITEBYTECODE": "1",
        }
    )
    results = []
    for config in selected:
        config_knobs = knobs(config, paths, contract, args.output_root / "cache", args.prefix_caching)
        for trace in (trace_by_tp[config.tp],):
            run_dir = args.output_root / "runs" / config.name / trace["label"]
            result_path = run_dir / "result.json"
            if args.resume and result_path.is_file():
                previous = json.loads(result_path.read_text())
                if (
                    previous.get("status") == "completed"
                    and previous.get("allreduce_csv_sha256")
                    == allreduce_csv_sha256
                ):
                    results.append(previous)
                    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"qwen235_v020_{config.name}_{trace['label']}",
                knobs=config_knobs,
            )
            command.extend(["--cudagraph_capture_sizes", *(str(v) for v in contract[config.name]["capture_sizes"])])
            command = BASE.configure_cc_command(
                command,
                backend="vidur",
                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,
                )
            result = {
                "config": asdict(config) | {"name": config.name},
                "trace": Q30.trace_manifest_entry(trace),
                "elapsed_seconds": time.time() - started,
                "returncode": completed.returncode,
                "allreduce_csv_sha256": allreduce_csv_sha256,
            }
            fallback_evidence = Q30.collective_fallback_evidence(run_dir)
            if completed.returncode == 0 and fallback_evidence:
                result.update(
                    status="failed",
                    failure_class="collective_profile_fallback",
                    collective_fallback_evidence=fallback_evidence,
                )
            elif completed.returncode == 0:
                metrics = Q30.find_request_metrics(run_dir)
                result.update(status="completed", metrics=Q30.score(metrics, trace["shapes"]), request_metrics_sha256=BASE.sha256(metrics))
            else:
                stderr_text = (run_dir / "stderr.log").read_text(errors="replace")
                result.update(status="failed", failure_class=Q30.classify_frontier_failure(stderr_text))
            BASE.write_json(result_path, result)
            results.append(result)
            print(json.dumps({"config": config.name, "trace": trace["label"], "status": result["status"]}, sort_keys=True), flush=True)

    manifest = {
        "schema": "qwen235-v020-frontier-latency-surface-v1",
        "frontier_commit": subprocess.check_output(["git", "-C", str(args.frontier_source), "rev-parse", "HEAD"], text=True).strip(),
        "profiles": {name: BASE.sha256(path) for name, path in paths.items()},
        "runtime_contract_sha256": BASE.sha256(args.runtime_contract),
        "collective": {
            "backend": "vidur",
            "allreduce_csv": str(args.allreduce_csv),
            "allreduce_csv_sha256": allreduce_csv_sha256,
            "measured_profile_validation": allreduce_profile,
        },
        "prefix_caching": args.prefix_caching,
        "results": results,
    }
    BASE.write_json(args.output_root / "frontier_surface.json", manifest)


if __name__ == "__main__":
    main()
