#!/usr/bin/env python3
"""Profile Qwen3's replicated MoE gate and fused top-k in vLLM 0.20."""

from __future__ import annotations

import argparse
import json
import socket
import statistics
import subprocess
from pathlib import Path
from typing import Any, Callable

import torch
import vllm


VLLM_VERSION = "0.20.0"
VLLM_COMMIT = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1"
HIDDEN_DIM = 2048
NUM_EXPERTS = 128
TOP_K = 8


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--vllm-source", type=Path, required=True)
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--num-tokens", type=int, nargs="+", required=True)
    parser.add_argument("--warmup-iters", type=int, default=5)
    parser.add_argument("--repeats", type=int, default=20)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument(
        "--profile-method",
        choices=("cuda_event", "record_function"),
        default="cuda_event",
    )
    parser.add_argument("--frontier-source", type=Path)
    return parser.parse_args()


def git_head(repo: Path) -> str:
    return subprocess.check_output(
        ["git", "-C", str(repo), "rev-parse", "HEAD"], text=True
    ).strip()


def stats_ms(samples: list[float]) -> dict[str, float]:
    return {
        "min": min(samples),
        "max": max(samples),
        "mean": statistics.fmean(samples),
        "median": statistics.median(samples),
        "std": statistics.pstdev(samples),
    }


def measure_ms(
    fn: Callable[[], Any],
    warmup_iters: int,
    repeats: int,
    *,
    profile_method: str,
    trace_root: Path | None = None,
    operation_name: str | None = None,
    record_function_tracer: type | None = None,
) -> tuple[Any, dict[str, float]]:
    result = None
    for _ in range(warmup_iters):
        result = fn()
    torch.accelerator.synchronize()
    if profile_method == "record_function":
        if trace_root is None or operation_name is None or record_function_tracer is None:
            raise ValueError("record_function profiling requires tracer metadata")
        tracer = record_function_tracer(str(trace_root))
        with tracer:
            for _ in range(repeats):
                with torch.profiler.record_function(f"vidur_{operation_name}"):
                    result = fn()
        stats = tracer.get_operation_time_stats()
        if operation_name not in stats:
            raise RuntimeError(f"missing RecordFunctionTracer stats for {operation_name}")
        return result, {
            name: float(value) for name, value in stats[operation_name].items()
        }

    samples: list[float] = []
    for _ in range(repeats):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        result = fn()
        end.record()
        torch.accelerator.synchronize()
        samples.append(float(start.elapsed_time(end)))
    return result, stats_ms(samples)


def main() -> None:
    args = parse_args()
    if vllm.__version__ != VLLM_VERSION:
        raise SystemExit(f"expected vLLM {VLLM_VERSION}, got {vllm.__version__}")
    source_head = git_head(args.vllm_source)
    if source_head != VLLM_COMMIT:
        raise SystemExit(f"expected vLLM source {VLLM_COMMIT}, got {source_head}")
    if args.profile_method == "record_function" and args.frontier_source is None:
        raise SystemExit("--frontier-source is required for --profile-method record_function")

    raw_model_config = json.loads(args.model.joinpath("config.json").read_text())
    observed = {
        "hidden_size": raw_model_config.get("hidden_size"),
        "num_experts": raw_model_config.get("num_experts"),
        "num_experts_per_tok": raw_model_config.get("num_experts_per_tok"),
        "norm_topk_prob": raw_model_config.get("norm_topk_prob"),
    }
    expected = {
        "hidden_size": HIDDEN_DIM,
        "num_experts": NUM_EXPERTS,
        "num_experts_per_tok": TOP_K,
        "norm_topk_prob": True,
    }
    if observed != expected:
        raise SystemExit(f"model contract mismatch: expected {expected}, got {observed}")

    from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
    from vllm.distributed import (
        destroy_distributed_environment,
        destroy_model_parallel,
        init_distributed_environment,
        initialize_model_parallel,
    )
    from vllm.model_executor.layers.fused_moe import fused_topk
    from vllm.model_executor.layers.linear import ReplicatedLinear

    record_function_tracer = None
    if args.profile_method == "record_function":
        import sys

        sys.path.insert(0, str(args.frontier_source.resolve()))
        from frontier.profiling.utils.record_function_tracer import RecordFunctionTracer

        record_function_tracer = RecordFunctionTracer

    device = torch.device(args.device)
    torch.accelerator.set_device_index(device)
    torch.manual_seed(20260716)
    model_config = ModelConfig(
        model=str(args.model),
        dtype="bfloat16",
        max_model_len=8192,
        skip_tokenizer_init=True,
        generation_config="vllm",
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    if args.profile_method == "record_function":
        (args.output.parent / "profiler_traces").mkdir(exist_ok=True)

    rows: list[dict[str, Any]] = []
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
        listener.bind(("127.0.0.1", 0))
        distributed_init_method = f"tcp://127.0.0.1:{listener.getsockname()[1]}"
    init_distributed_environment(
        world_size=1,
        rank=0,
        local_rank=0,
        distributed_init_method=distributed_init_method,
    )
    try:
        with set_current_vllm_config(VllmConfig(model_config=model_config)):
            initialize_model_parallel(tensor_model_parallel_size=1)
            gate = ReplicatedLinear(
                HIDDEN_DIM,
                NUM_EXPERTS,
                bias=False,
                quant_config=None,
                prefix="model.layers.0.mlp.gate",
            ).to(device=device, dtype=torch.bfloat16)
            gate.weight.data.uniform_(-0.01, 0.01)

            for num_tokens in args.num_tokens:
                hidden = torch.empty(
                    (num_tokens, HIDDEN_DIM), device=device, dtype=torch.bfloat16
                ).uniform_(-0.1, 0.1)
                logits, gate_time = measure_ms(
                    lambda: gate(hidden)[0],
                    args.warmup_iters,
                    args.repeats,
                    profile_method=args.profile_method,
                    trace_root=args.output.parent,
                    operation_name="moe_gating_linear",
                    record_function_tracer=record_function_tracer,
                )
                topk_result, topk_time = measure_ms(
                    lambda: fused_topk(hidden, logits, TOP_K, renormalize=True),
                    args.warmup_iters,
                    args.repeats,
                    profile_method=args.profile_method,
                    trace_root=args.output.parent,
                    operation_name="moe_gating_routing_topk",
                    record_function_tracer=record_function_tracer,
                )

                def gate_and_topk() -> tuple[
                    torch.Tensor, torch.Tensor, torch.Tensor
                ]:
                    current_logits, _ = gate(hidden)
                    return fused_topk(
                        hidden, current_logits, TOP_K, renormalize=True
                    )

                combined_result, combined_time = measure_ms(
                    gate_and_topk,
                    args.warmup_iters,
                    args.repeats,
                    profile_method=args.profile_method,
                    trace_root=args.output.parent,
                    operation_name="moe_gating_linear_and_routing_topk",
                    record_function_tracer=record_function_tracer,
                )
                topk_weights, topk_ids, _ = topk_result
                combined_weights, combined_ids, _ = combined_result
                if logits.shape != (num_tokens, NUM_EXPERTS):
                    raise SystemExit(
                        f"invalid gate output shape: {tuple(logits.shape)}"
                    )
                if topk_ids.shape != (num_tokens, TOP_K):
                    raise SystemExit(f"invalid top-k shape: {tuple(topk_ids.shape)}")
                torch.testing.assert_close(
                    topk_weights.sum(dim=-1),
                    torch.ones(num_tokens, device=device),
                    atol=1e-5,
                    rtol=1e-5,
                )
                torch.testing.assert_close(combined_weights, topk_weights)
                torch.testing.assert_close(combined_ids, topk_ids)
                additive_median = gate_time["median"] + topk_time["median"]
                row = {
                    "num_tokens": num_tokens,
                    "gate_linear_time_ms": gate_time,
                    "routing_topk_time_ms": topk_time,
                    "gate_plus_topk_time_ms": combined_time,
                    "median_nonadditivity_ratio": (
                        combined_time["median"] / additive_median
                        if additive_median > 0
                        else 1.0
                    ),
                }
                rows.append(row)
                print(json.dumps(row, sort_keys=True), flush=True)
    finally:
        destroy_model_parallel()
        destroy_distributed_environment()

    payload = {
        "schema_version": "qwen30_vllm020_router_raw.v1",
        "environment": {
            "vllm_version": vllm.__version__,
            "vllm_source_commit": source_head,
            "torch_version": torch.__version__,
            "torch_cuda": torch.version.cuda,
            "gpu": torch.cuda.get_device_name(device),
            "model": str(args.model),
            "dtype": "bfloat16",
            "gate_replication": "replicated_across_tp",
            "top_k": TOP_K,
            "norm_topk_prob": True,
            "profile_method": args.profile_method,
        },
        "measurement_scope": (
            "vLLM ReplicatedLinear gate and fused_topk; measured separately and "
            "as the actual sequential router path"
        ),
        "rows": rows,
    }
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")


if __name__ == "__main__":
    main()
