#!/usr/bin/env python3
"""Run Frontier on the exact single/concurrency-2 T0 smoke workload."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import os
import subprocess
import time
from pathlib import Path
from typing import Any


MODEL = "Qwen3-235B-A22B-FP8"


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, default=Path("/usr/bin/python3.12"))
    parser.add_argument("--output-root", type=Path, required=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)
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")


def write_trace(path: Path, request_count: int) -> 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 _ in range(request_count):
            writer.writerow(
                {
                    "arrived_at": 0.0,
                    "num_prefill_tokens": 2048,
                    "num_decode_tokens": 128,
                    "slo_ttft_ms": 1256.0,
                }
            )


def build_command(
    *,
    args: argparse.Namespace,
    trace: Path,
    run_dir: Path,
    run_id: str,
) -> list[str]:
    compute_root = args.profile_root / "compute/h20" / MODEL
    network = args.profile_root / "network/h20_nccl/all_reduce.csv"
    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",
        "4",
        "--replica_config_attn_data_parallel_size",
        "1",
        "--replica_config_moe_tensor_parallel_size",
        "4",
        "--replica_config_moe_expert_parallel_size",
        "1",
        "--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(run_dir / "cache/collectives"),
        "--vidur_cc_backend_config_all_reduce_input_file",
        str(network),
        "--replica_scheduler_config_type",
        "vllm_v1",
        "--decode_cuda_graph_mode",
        "none",
        "--vllm_v1_scheduler_config_batch_size_cap",
        "64",
        "--vllm_v1_scheduler_config_block_size",
        "16",
        "--vllm_v1_scheduler_config_num_blocks",
        "26101",
        "--vllm_v1_scheduler_config_num_blocks_mode",
        "explicit",
        "--vllm_v1_scheduler_config_max_tokens_in_batch",
        "8192",
        "--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(compute_root / "linear_op.csv"),
        "--random_forrest_execution_time_predictor_config_atten_input_file",
        str(compute_root / "attention.csv"),
        "--random_forrest_execution_time_predictor_config_moe_input_file",
        str(compute_root / "moe.csv"),
        "--random_forrest_execution_time_predictor_config_all_reduce_input_file",
        str(network),
        "--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(run_dir / "cache/execution"),
        "--metrics_config_output_dir",
        str(run_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 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, found {matches}")
    return matches[0]


def score(metrics_path: Path) -> dict[str, Any]:
    with metrics_path.open(newline="") as source:
        rows = list(csv.DictReader(source))
    requests = []
    for row in rows:
        ttft_ms = float(row["ttft"])
        e2e_ms = float(row["request_e2e_time"])
        decode_tokens = int(float(row["request_num_decode_tokens"]))
        aligned_tpot_ms = (
            (e2e_ms - ttft_ms) / (decode_tokens - 1) if decode_tokens > 1 else 0.0
        )
        requests.append(
            {
                "request_id": int(row["Request Id"]),
                "prompt_tokens": int(float(row["request_num_prefill_tokens"])),
                "completion_tokens": decode_tokens,
                "ttft_ms": ttft_ms,
                "e2e_ms": e2e_ms,
                "tpot_ms_aligned": aligned_tpot_ms,
                "frontier_decode_e2e_time_per_token_ms": float(
                    row["decode_e2e_time_per_token"]
                ),
                "joint_slo_pass": ttft_ms <= 1256.0 and aligned_tpot_ms <= 40.0,
            }
        )
    return {
        "request_count": len(requests),
        "joint_slo_pass_count": sum(row["joint_slo_pass"] for row in requests),
        "requests": requests,
    }


def main() -> None:
    args = parse_args()
    source = args.frontier_source.resolve()
    profiles = args.profile_root.resolve()
    output = args.output_root.resolve()
    required = [
        source / "pyproject.toml",
        profiles / f"compute/h20/{MODEL}/linear_op.csv",
        profiles / f"compute/h20/{MODEL}/attention.csv",
        profiles / f"compute/h20/{MODEL}/moe.csv",
        profiles / "network/h20_nccl/all_reduce.csv",
    ]
    missing = [str(path) for path in required if not path.is_file()]
    if missing:
        raise FileNotFoundError(missing)

    results: dict[str, Any] = {}
    for label, count in (("single", 1), ("concurrency2", 2)):
        run_dir = output / label
        trace = run_dir / "trace.csv"
        write_trace(trace, count)
        command = build_command(
            args=args,
            trace=trace,
            run_dir=run_dir,
            run_id=f"t0_{label}",
        )
        write_json(run_dir / "command.json", command)
        env = dict(os.environ)
        env["PYTHONPATH"] = str(source)
        started = time.time()
        with (run_dir / "stdout.log").open("w") as stdout:
            completed = subprocess.run(
                command,
                cwd=source,
                env=env,
                stdout=stdout,
                stderr=subprocess.STDOUT,
                check=False,
                text=True,
                timeout=300,
            )
        if completed.returncode != 0:
            raise RuntimeError(
                f"Frontier {label} failed with {completed.returncode}; "
                f"see {run_dir / 'stdout.log'}"
            )
        metrics = find_request_metrics(run_dir)
        results[label] = {
            "elapsed_seconds": time.time() - started,
            "trace_sha256": sha256(trace),
            "request_metrics_path": str(metrics),
            "request_metrics_sha256": sha256(metrics),
            **score(metrics),
        }

    profile_files = required[1:]
    payload = {
        "schema": "frontier-qwen235b-t0-smoke-v1",
        "contract": {
            "topology": "TP4/DP1/MoE-TP4/EP1",
            "mns": 64,
            "mbt": 8192,
            "block_size": 16,
            "num_gpu_blocks": 26101,
            "prefix_caching": False,
            "input_tokens": 2048,
            "output_tokens": 128,
            "arrivals": "all at t=0",
        },
        "frontier": {
            "source": str(source),
            "declared_upstream_commit": "d9cfeb6d8791fbf2f295dd9744c56a666171776e",
            "python_and_config_tree_sha256": "172fc7ae19c40e67030208ff488d0d5d90764888ed76a7a543be6037ba62dc11",
        },
        "profiles": {str(path): sha256(path) for path in profile_files},
        "results": results,
    }
    write_json(output / "summary.json", payload)
    print(json.dumps(payload["results"], indent=2, sort_keys=True))


if __name__ == "__main__":
    main()