#!/usr/bin/env python3
"""Freeze a profile-only Frontier ranking for the Qwen235B prefill grid."""

from __future__ import annotations

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


SCHEMA = "frontier-qwen235b-prefill-grid-v1"
MODEL = "Qwen3-235B-A22B-FP8"
TRACE_SHA256 = "f878e9af18f94dcfaced94a8e1e6b20a2f7d97d64aa862448025660dbbd965b2"
PROFILE_RELATIVE = Path("compute/h20") / MODEL
NETWORK_RELATIVE = Path("network/h20_nccl/all_reduce.csv")
SEARCH_LOW = 0.0
SEARCH_HIGH = 0.125
SEARCH_PROBES = 6
WINDOW_DURATION_S = 600.0
MAX_INPUT_TOKENS = 32768
MAX_MODEL_TOKENS = 40960
BLOCK_SIZE_TOKENS = 16
MOE_ROUTING_SEED = 42


@dataclass(frozen=True)
class GridConfig:
    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}"

    @property
    def gpu_count(self) -> int:
        return self.tp


GRID = tuple(
    GridConfig(
        tp=tp,
        mns=mns,
        mbt=mbt,
        moe_tp=4 if tp == 4 else 1,
        moe_ep=1 if tp == 4 else 8,
        # Measured by community vLLM on the same checkpoint/runtime. The real
        # grid will use --num-gpu-blocks-override with these same values.
        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 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 order_hash(values: Iterable[object]) -> str:
    payload = "\n".join(str(value) for value in values).encode()
    return hashlib.sha256(payload).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 sha256_bytes(payload: bytes) -> str:
    return hashlib.sha256(payload).hexdigest()


def frontier_source_fingerprint(
    source: Path, declared_commit: str
) -> dict[str, Any]:
    model_config = source / "data" / "config" / "models" / f"{MODEL}.json"
    if not model_config.is_file():
        raise FileNotFoundError(f"Frontier model config is missing: {model_config}")
    device_config = source / "data" / "config" / "device" / "h20.json"
    files = sorted((source / "frontier").glob("**/*.py")) + [
        model_config,
        device_config,
    ]
    missing = [str(path) for path in files if not path.is_file()]
    if missing:
        raise FileNotFoundError(f"Frontier fingerprint inputs are missing: {missing}")
    digest = hashlib.sha256()
    for path in files:
        relative = path.relative_to(source).as_posix()
        digest.update(relative.encode() + b"\0")
        digest.update(bytes.fromhex(sha256(path)))
    fingerprint = {
        "declared_upstream_commit": declared_commit,
        "python_and_config_tree_sha256": digest.hexdigest(),
        "fingerprinted_file_count": len(files),
        "model_config": {
            "path": str(model_config.resolve()),
            "sha256": sha256(model_config),
        },
    }
    git_dir = source / ".git"
    if git_dir.exists():
        commit = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            cwd=source,
            check=True,
            stdout=subprocess.PIPE,
            text=True,
        ).stdout.strip()
        status = subprocess.run(
            ["git", "status", "--porcelain=v1"],
            cwd=source,
            check=True,
            stdout=subprocess.PIPE,
            text=True,
        ).stdout
        diff = subprocess.run(
            ["git", "diff", "--binary", "HEAD"],
            cwd=source,
            check=True,
            stdout=subprocess.PIPE,
        ).stdout
        fingerprint.update(
            {
                "git_commit": commit,
                "status_porcelain": status.splitlines(),
                "tracked_diff_sha256": sha256_bytes(diff),
            }
        )
    else:
        fingerprint["git_metadata"] = "absent_source_snapshot"
    return fingerprint


def anchor_key(anchor: float) -> str:
    return f"{anchor:.12f}"


def anchor_filename(anchor: float) -> str:
    return f"u_{anchor_key(anchor).replace('.', 'p')}.csv"


def binary_search_lattice() -> list[float]:
    intervals = [(SEARCH_LOW, SEARCH_HIGH)]
    anchors: list[float] = []
    for _ in range(SEARCH_PROBES):
        next_intervals = []
        for low, high in intervals:
            midpoint = (low + high) / 2.0
            anchors.append(midpoint)
            next_intervals.extend(((low, midpoint), (midpoint, high)))
        intervals = next_intervals
    return sorted(set(anchors))


def ttft_slo_ms(input_tokens: int) -> float:
    return 1000.0 if input_tokens <= 8191 else 2000.0


def load_trace_rows(path: Path) -> list[dict[str, Any]]:
    rows = []
    with path.open(encoding="utf-8") as source:
        for source_index, line in enumerate(source):
            if not line.strip():
                continue
            raw = json.loads(line)
            input_tokens = int(raw["input_length"])
            if not 0 <= input_tokens <= MAX_INPUT_TOKENS:
                continue
            timestamp = float(raw["timestamp"])
            sampling_u = float(raw["sampling_u"])
            rows.append(
                {
                    "arrived_at": timestamp,
                    "num_prefill_tokens": input_tokens,
                    "num_decode_tokens": 1,
                    "source_row_index": source_index,
                    "source_request_id": str(
                        raw.get("request_id") or raw.get("id") or source_index
                    ),
                    "sampling_u": sampling_u,
                    "slo_ttft_ms": ttft_slo_ms(input_tokens),
                }
            )
    # This reproduces AITuner's stable arrival-only sort. Tied timestamps retain
    # source order; sampling_u must not become a secondary key.
    rows.sort(key=lambda row: row["arrived_at"])
    return rows


def prepare_traces(trace: Path, output_root: Path, expected_sha256: str) -> Path:
    actual_sha256 = sha256(trace)
    if expected_sha256 and actual_sha256 != expected_sha256:
        raise ValueError(
            f"trace SHA256 mismatch: expected={expected_sha256}, actual={actual_sha256}"
        )
    rows = load_trace_rows(trace)
    if not rows:
        raise ValueError("no trace rows remain after input-length filtering")
    if any(
        rows[index]["arrived_at"] > rows[index + 1]["arrived_at"]
        for index in range(len(rows) - 1)
    ):
        raise ValueError("materialized trace is not monotonic by arrival")

    trace_dir = output_root / "traces"
    trace_dir.mkdir(parents=True, exist_ok=True)
    fields = [
        "arrived_at",
        "num_prefill_tokens",
        "num_decode_tokens",
        "source_row_index",
        "source_request_id",
        "sampling_u",
        "slo_ttft_ms",
    ]
    anchors: dict[str, Any] = {}
    for anchor in binary_search_lattice():
        selected = [row for row in rows if row["sampling_u"] <= anchor]
        target = trace_dir / anchor_filename(anchor)
        with target.open("w", encoding="utf-8", newline="") as output:
            writer = csv.DictWriter(output, fieldnames=fields)
            writer.writeheader()
            writer.writerows(selected)
        key = anchor_key(anchor)
        anchors[key] = {
            "anchor": anchor,
            "path": str(target.resolve()),
            "sha256": sha256(target),
            "request_count": len(selected),
            "request_rate": len(selected) / WINDOW_DURATION_S,
            "source_row_order_sha256": order_hash(
                row["source_row_index"] for row in selected
            ),
            "arrival_order_sha256": order_hash(
                f"{row['arrived_at']:.12f}" for row in selected
            ),
            "input_length_order_sha256": order_hash(
                row["num_prefill_tokens"] for row in selected
            ),
        }

    manifest = {
        "schema": SCHEMA,
        "source": {
            "path": str(trace.resolve()),
            "sha256": actual_sha256,
            "filtered_request_count": len(rows),
        },
        "selection_contract": {
            "input_tokens": [0, MAX_INPUT_TOKENS],
            "completion_tokens_override": 1,
            "stable_sort_key": "arrived_at_only",
            "selection": "sampling_u <= anchor",
            "window_duration_s": WINDOW_DURATION_S,
            "search_low": SEARCH_LOW,
            "search_high": SEARCH_HIGH,
            "search_probes": SEARCH_PROBES,
            "lattice_anchor_count": len(anchors),
            "ttft_slo_ms": {"input_le_8191": 1000, "otherwise": 2000},
            "target_pass_rate": 0.95,
        },
        "anchors": anchors,
    }
    manifest_path = output_root / "trace_manifest.json"
    write_json(manifest_path, manifest)
    return manifest_path


def resolve_profile_paths(profile_root: Path) -> dict[str, Path]:
    paths = {
        "linear": profile_root / PROFILE_RELATIVE / "linear_op.csv",
        "attention": profile_root / PROFILE_RELATIVE / "attention.csv",
        "moe": profile_root / PROFILE_RELATIVE / "moe.csv",
        "all_reduce": profile_root / NETWORK_RELATIVE,
        "manifest": profile_root / "profile_manifest.json",
    }
    missing = [f"{name}={path}" for name, path in paths.items() if not path.is_file()]
    if missing:
        raise FileNotFoundError("missing profile inputs: " + ", ".join(missing))
    return paths


def build_command(
    *,
    python: Path,
    frontier_source: Path,
    profile_root: Path,
    profile_paths: dict[str, Path],
    trace: Path,
    config: GridConfig,
    probe_dir: Path,
    run_id: str,
    cache_root: Path,
) -> list[str]:
    return [
        str(python),
        "-m",
        "frontier.main",
        "--simulation_mode",
        "offline",
        # Frontier's offline default rewrites every request arrival to time 0.
        # Preserve the trace arrival process so the simulator and serving run
        # exercise the same queueing workload.
        "--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",
        str(MOE_ROUTING_SEED),
        "--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(profile_root),
        "--vidur_cc_backend_config_cache_dir",
        str(cache_root / "collectives"),
        "--vidur_cc_backend_config_all_reduce_input_file",
        str(profile_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",
        str(BLOCK_SIZE_TOKENS),
        "--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",
        str(MAX_MODEL_TOKENS),
        "--no-random_forrest_execution_time_predictor_config_enable_dummy_mode",
        "--random_forrest_execution_time_predictor_config_linear_op_input_file",
        str(profile_paths["linear"]),
        "--random_forrest_execution_time_predictor_config_atten_input_file",
        str(profile_paths["attention"]),
        "--random_forrest_execution_time_predictor_config_moe_input_file",
        str(profile_paths["moe"]),
        "--random_forrest_execution_time_predictor_config_all_reduce_input_file",
        str(profile_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",
        str(MAX_INPUT_TOKENS + 1),
        "--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_root / "execution"),
        "--metrics_config_output_dir",
        str(probe_dir / "metrics"),
        "--metrics_config_run_id",
        run_id,
        "--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 score_request_metrics(trace: Path, request_metrics: Path) -> dict[str, Any]:
    with trace.open(encoding="utf-8", newline="") as source:
        trace_rows = list(csv.DictReader(source))
    with request_metrics.open(encoding="utf-8", newline="") as source:
        metric_rows = list(csv.DictReader(source))
    if len(metric_rows) != len(trace_rows):
        raise ValueError(
            f"request count mismatch: trace={len(trace_rows)}, metrics={len(metric_rows)}"
        )
    metrics_by_id = {int(row["Request Id"]): row for row in metric_rows}
    expected_ids = set(range(len(trace_rows)))
    if set(metrics_by_id) != expected_ids:
        raise ValueError("Frontier Request Ids do not match trace row positions")

    passed = 0
    ttfts = []
    for request_id, trace_row in enumerate(trace_rows):
        ttft = float(metrics_by_id[request_id]["ttft"])
        threshold = float(trace_row["slo_ttft_ms"])
        ttfts.append(ttft)
        passed += int(ttft <= threshold)
    count = len(trace_rows)
    pass_rate = passed / count if count else 0.0
    ordered = sorted(ttfts)

    def percentile(fraction: float) -> float | None:
        if not ordered:
            return None
        index = round((len(ordered) - 1) * fraction)
        return ordered[index]

    return {
        "request_count": count,
        "passed_request_count": passed,
        "slo_pass_rate": pass_rate,
        "feasible": pass_rate >= 0.95,
        "ttft_ms": {
            "min": min(ordered) if ordered else None,
            "p50": percentile(0.50),
            "p95": percentile(0.95),
            "p99": percentile(0.99),
            "max": max(ordered) if ordered else None,
        },
    }


def find_request_metrics(probe_dir: Path) -> Path:
    candidates = list((probe_dir / "metrics").glob("**/request_metrics.csv"))
    if len(candidates) != 1:
        raise ValueError(
            f"expected one request_metrics.csv under {probe_dir}, got {candidates}"
        )
    return candidates[0]


def run_probe(
    *,
    python: Path,
    frontier_source: Path,
    profile_root: Path,
    profile_paths: dict[str, Path],
    trace_record: dict[str, Any],
    config: GridConfig,
    probe_index: int,
    output_root: Path,
    cache_root: Path,
) -> dict[str, Any]:
    anchor = float(trace_record["anchor"])
    probe_dir = output_root / "runs" / config.name / f"probe_{probe_index}_{anchor_filename(anchor)[:-4]}"
    result_path = probe_dir / "result.json"
    if result_path.is_file():
        result = json.loads(result_path.read_text())
        if result.get("status") == "completed":
            return result
    probe_dir.mkdir(parents=True, exist_ok=True)
    trace = Path(trace_record["path"])
    run_id = f"{config.name}_probe{probe_index}_{anchor_filename(anchor)[:-4]}"
    command = build_command(
        python=python,
        frontier_source=frontier_source,
        profile_root=profile_root,
        profile_paths=profile_paths,
        trace=trace,
        config=config,
        probe_dir=probe_dir,
        run_id=run_id,
        cache_root=cache_root,
    )
    write_json(probe_dir / "command.json", command)
    environment = os.environ.copy()
    environment.update(
        {
            "PYTHONPATH": str(frontier_source),
            "WANDB_DISABLED": "true",
            "VIDUR_DISABLE_WANDB": "1",
        }
    )
    started = time.time()
    with (probe_dir / "stdout.log").open("w", encoding="utf-8") as output:
        completed = subprocess.run(
            command,
            cwd=frontier_source,
            env=environment,
            stdout=output,
            stderr=subprocess.STDOUT,
            check=False,
        )
    elapsed = time.time() - started
    if completed.returncode != 0:
        result = {
            "status": "failed",
            "returncode": completed.returncode,
            "elapsed_seconds": elapsed,
            "config": asdict(config),
            "anchor": anchor,
            "trace_sha256": trace_record["sha256"],
        }
        write_json(result_path, result)
        raise RuntimeError(f"Frontier probe failed: {config.name}, anchor={anchor}")

    request_metrics = find_request_metrics(probe_dir)
    score = score_request_metrics(trace, request_metrics)
    result = {
        "status": "completed",
        "elapsed_seconds": elapsed,
        "config": asdict(config),
        "anchor": anchor,
        "trace_path": str(trace),
        "trace_sha256": trace_record["sha256"],
        "request_metrics_path": str(request_metrics),
        "request_metrics_sha256": sha256(request_metrics),
        "request_rate": score["request_count"] / WINDOW_DURATION_S,
        "request_rate_per_gpu": score["request_count"]
        / WINDOW_DURATION_S
        / config.gpu_count,
        **score,
    }
    write_json(result_path, result)
    return result


def selected_configs(names: list[str] | None) -> list[GridConfig]:
    if not names:
        return list(GRID)
    by_name = {config.name: config for config in GRID}
    unknown = sorted(set(names) - set(by_name))
    if unknown:
        raise ValueError(f"unknown configs: {unknown}; available={sorted(by_name)}")
    return [by_name[name] for name in names]


def run_grid(args: argparse.Namespace) -> None:
    trace_manifest = json.loads(args.trace_manifest.read_text())
    if trace_manifest.get("schema") != SCHEMA:
        raise ValueError(f"unexpected trace manifest schema: {trace_manifest.get('schema')}")
    profile_paths = resolve_profile_paths(args.profile_root)
    args.output_root.mkdir(parents=True, exist_ok=True)
    cache_root = args.output_root / "cache"
    configs = selected_configs(args.config)
    run_manifest = {
        "schema": SCHEMA,
        "frontier": {
            "source": str(args.frontier_source.resolve()),
            "python": str(args.python.resolve()),
            "fingerprint": frontier_source_fingerprint(
                args.frontier_source, args.frontier_commit
            ),
        },
        "trace_manifest": {
            "path": str(args.trace_manifest.resolve()),
            "sha256": sha256(args.trace_manifest),
        },
        "profiles": {
            name: {"path": str(path.resolve()), "sha256": sha256(path)}
            for name, path in profile_paths.items()
        },
        "kv_capacity_evidence": {
            "tp4": {
                "path": str(args.tp4_capacity_artifact.resolve()),
                "sha256": sha256(args.tp4_capacity_artifact),
                "num_gpu_blocks": 26101,
            },
            "tp8": {
                "path": str(args.tp8_capacity_artifact.resolve()),
                "sha256": sha256(args.tp8_capacity_artifact),
                "num_gpu_blocks": 62351,
            },
        },
        "contract": {
            "metric": "maximum SLO-feasible request_rate_per_gpu",
            "target_pass_rate": 0.95,
            "cpu_overhead_modeling": "skipped_no_community_vllm_native_records",
            "end_to_end_calibration": False,
            "moe_routing_mode": "simulation",
            "moe_routing_seed": MOE_ROUTING_SEED,
            "kv_blocks_source": "community_vllm_measured_and_fixed_per_topology",
            "network_device": "h20_dgx",
        },
        "configs": [asdict(config) | {"name": config.name} for config in configs],
    }
    write_json(args.output_root / "run_manifest.json", run_manifest)

    all_results = {}
    for config in configs:
        low = SEARCH_LOW
        high = SEARCH_HIGH
        probes = []
        best = None
        for probe_index in range(SEARCH_PROBES):
            anchor = (low + high) / 2.0
            record = trace_manifest["anchors"].get(anchor_key(anchor))
            if record is None:
                raise ValueError(f"trace manifest lacks anchor {anchor_key(anchor)}")
            result = run_probe(
                python=args.python,
                frontier_source=args.frontier_source,
                profile_root=args.profile_root,
                profile_paths=profile_paths,
                trace_record=record,
                config=config,
                probe_index=probe_index,
                output_root=args.output_root,
                cache_root=cache_root,
            )
            probes.append(result)
            print(
                json.dumps(
                    {
                        "config": config.name,
                        "probe": probe_index,
                        "anchor": anchor,
                        "pass_rate": result["slo_pass_rate"],
                        "feasible": result["feasible"],
                        "elapsed_seconds": result["elapsed_seconds"],
                    },
                    sort_keys=True,
                ),
                flush=True,
            )
            if result["feasible"]:
                low = anchor
                best = result
            else:
                high = anchor
        summary = {
            "config": asdict(config) | {"name": config.name},
            "probes": probes,
            "capacity_interval": [low, high],
            "best_feasible": best,
        }
        write_json(args.output_root / "results" / f"{config.name}.json", summary)
        all_results[config.name] = summary

    ranked = sorted(
        all_results.values(),
        key=lambda item: (
            -(
                item["best_feasible"]["request_rate_per_gpu"]
                if item["best_feasible"] is not None
                else -1.0
            ),
            item["config"]["name"],
        ),
    )
    freeze = {
        "schema": SCHEMA,
        "run_manifest_sha256": sha256(args.output_root / "run_manifest.json"),
        "ranking": [
            {
                "rank": index + 1,
                "config": item["config"],
                "capacity_interval_sampling_u": item["capacity_interval"],
                "best_feasible": item["best_feasible"],
            }
            for index, item in enumerate(ranked)
        ],
    }
    write_json(args.output_root / "frontier_ranking_frozen.json", freeze)
    print(json.dumps(freeze, indent=2), flush=True)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)
    prepare = subparsers.add_parser("prepare")
    prepare.add_argument("--trace", type=Path, required=True)
    prepare.add_argument("--output-root", type=Path, required=True)
    prepare.add_argument("--expected-trace-sha256", default=TRACE_SHA256)

    run = subparsers.add_parser("run")
    run.add_argument("--frontier-source", type=Path, required=True)
    run.add_argument(
        "--frontier-commit",
        default="d9cfeb6d8791fbf2f295dd9744c56a666171776e",
    )
    run.add_argument("--python", type=Path, required=True)
    run.add_argument("--profile-root", type=Path, required=True)
    run.add_argument("--trace-manifest", type=Path, required=True)
    run.add_argument("--output-root", type=Path, required=True)
    run.add_argument("--tp4-capacity-artifact", type=Path, required=True)
    run.add_argument("--tp8-capacity-artifact", type=Path, required=True)
    run.add_argument("--config", action="append")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.command == "prepare":
        manifest = prepare_traces(
            args.trace, args.output_root, args.expected_trace_sha256
        )
        print(manifest)
        return
    if args.command == "run":
        run_grid(args)
        return
    raise AssertionError(args.command)


if __name__ == "__main__":
    main()
