#!/usr/bin/env python3
"""Capture exact Qwen3 routed-expert IDs from vLLM 0.20 on trace prompts."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import subprocess
from pathlib import Path
from typing import Any

import numpy as np
import torch
import vllm


VLLM_VERSION = "0.20.0"
VLLM_COMMIT = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1"
NUM_EXPERTS = 128
TOP_K = 8
NUM_LAYERS = 48


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("--fixture", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--routes", type=Path, required=True)
    parser.add_argument("--decode-override", type=int)
    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 sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def common_prefix(left: list[int], right: list[int]) -> int:
    count = 0
    for lhs, rhs in zip(left, right):
        if lhs != rhs:
            break
        count += 1
    return count


def distribution(counts: np.ndarray) -> dict[str, Any]:
    values = counts.astype(np.float64)
    total = float(values.sum())
    mean = float(values.mean())
    probabilities = values[values > 0] / total
    entropy = float(-(probabilities * np.log2(probabilities)).sum())
    variance = float(((values - mean) ** 2).mean())
    ordered = np.sort(values)
    gini = float(
        2.0 * np.dot(np.arange(1, len(values) + 1), ordered)
        / (len(values) * total)
        - (len(values) + 1) / len(values)
    )
    hottest = np.argsort(values)[-8:][::-1]
    return {
        "total_routed_tokens": int(total),
        "tokens_per_expert_mean": mean,
        "load_cv": math.sqrt(variance) / mean,
        "load_gini": gini,
        "load_entropy_bits": entropy,
        "min_load_ratio": float(values.min() / mean),
        "max_load_ratio": float(values.max() / mean),
        "expert_utilization": float(np.count_nonzero(values) / len(values)),
        "hottest_experts": [int(value) for value in hottest],
        "hottest_counts": [int(values[value]) for value in hottest],
        "counts": counts.astype(int).tolist(),
    }


def phase_summary(routes: list[np.ndarray]) -> dict[str, Any]:
    counts = np.zeros(NUM_EXPERTS, dtype=np.int64)
    per_layer = np.zeros((NUM_LAYERS, NUM_EXPERTS), dtype=np.int64)
    token_count = 0
    for route in routes:
        token_count += route.shape[0]
        counts += np.bincount(route.reshape(-1), minlength=NUM_EXPERTS)
        for layer in range(NUM_LAYERS):
            per_layer[layer] += np.bincount(
                route[:, layer, :].reshape(-1), minlength=NUM_EXPERTS
            )
    return {
        "token_count": token_count,
        "all_layers": distribution(counts),
        "per_layer": [distribution(row) for row in per_layer],
    }


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}")
    rows = [json.loads(line) for line in args.fixture.read_text().splitlines() if line]
    if not rows:
        raise SystemExit("empty routing fixture")
    requested_decode = [
        args.decode_override
        if args.decode_override is not None
        else int(row["output_length"])
        for row in rows
    ]
    if any(value <= 0 for value in requested_decode):
        raise SystemExit("all requested decode lengths must be positive")

    from vllm import LLM, SamplingParams

    llm = LLM(
        model=str(args.model),
        dtype="bfloat16",
        tensor_parallel_size=1,
        max_model_len=16384,
        max_num_batched_tokens=8192,
        max_num_seqs=64,
        gpu_memory_utilization=0.90,
        enable_chunked_prefill=True,
        enable_prefix_caching=True,
        enable_return_routed_experts=True,
        attention_backend="FLASH_ATTN",
        disable_log_stats=False,
    )
    sampling = [
        SamplingParams(temperature=0, min_tokens=value, max_tokens=value)
        for value in requested_decode
    ]
    conversations = [
        [{"role": "user", "content": row["prompt"]}] for row in rows
    ]
    outputs = llm.chat(conversations, sampling_params=sampling, use_tqdm=False)
    if len(outputs) != len(rows):
        raise SystemExit(f"expected {len(rows)} outputs, got {len(outputs)}")

    prompt_tokens_by_chat: dict[str, list[int]] = {}
    prefill_routes: list[np.ndarray] = []
    decode_routes: list[np.ndarray] = []
    raw_routes: dict[str, np.ndarray] = {}
    request_summaries = []
    for row, output, decode_tokens in zip(rows, outputs, requested_decode):
        completion = output.outputs[0]
        routed = completion.routed_experts
        if routed is None:
            raise SystemExit(f"row {row['row_id']} returned no routed experts")
        routed = np.asarray(routed)
        prompt_tokens = list(output.prompt_token_ids)
        generated_tokens = list(completion.token_ids)
        expected = len(prompt_tokens) + len(generated_tokens) - 1
        if routed.shape != (expected, NUM_LAYERS, TOP_K):
            raise SystemExit(
                f"row {row['row_id']} routes {routed.shape}, expected "
                f"{(expected, NUM_LAYERS, TOP_K)}"
            )
        if routed.min() < 0 or routed.max() >= NUM_EXPERTS:
            raise SystemExit(f"row {row['row_id']} returned invalid expert IDs")
        prefill = routed[: len(prompt_tokens)]
        decode = routed[len(prompt_tokens) :]
        if decode.shape[0] != decode_tokens - 1:
            raise SystemExit(f"row {row['row_id']} decode route length mismatch")
        prefill_routes.append(prefill)
        decode_routes.append(decode)
        raw_routes[f"row_{row['row_id']}"] = routed.astype(np.int16)
        prompt_tokens_by_chat[str(row["chat_id"])] = prompt_tokens
        request_summaries.append(
            {
                "fixture_index": row["fixture_index"],
                "row_id": row["row_id"],
                "turn": row["turn"],
                "input_length_trace": row["input_length"],
                "prompt_tokens_vllm": len(prompt_tokens),
                "chat_wrapper_delta": len(prompt_tokens) - int(row["input_length"]),
                "generated_tokens": len(generated_tokens),
                "requested_decode_tokens": decode_tokens,
                "routed_shape": list(routed.shape),
                "prompt_sha256": row["prompt_sha256"],
                "trace_hash_blocks": len(row["hash_ids"]),
            }
        )

    prefix_pairs = []
    by_chat = {str(row["chat_id"]): row for row in rows}
    for child in rows:
        parent = by_chat.get(str(child["parent_chat_id"]))
        if parent is None:
            continue
        parent_tokens = prompt_tokens_by_chat[str(parent["chat_id"])]
        child_tokens = prompt_tokens_by_chat[str(child["chat_id"])]
        prefix_pairs.append(
            {
                "parent_row_id": parent["row_id"],
                "child_row_id": child["row_id"],
                "trace_hash_common_prefix_blocks": common_prefix(
                    parent["hash_ids"], child["hash_ids"]
                ),
                "vllm_token_common_prefix": common_prefix(parent_tokens, child_tokens),
                "vllm_full_common_blocks_16": common_prefix(
                    parent_tokens, child_tokens
                )
                // 16,
            }
        )

    args.routes.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(args.routes, **raw_routes)
    payload = {
        "schema_version": "qwen30_vllm020_trace_routing.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(0),
            "model": str(args.model),
            "dtype": "bfloat16",
            "tensor_parallel_size": 1,
            "max_num_batched_tokens": 8192,
            "max_num_seqs": 64,
            "prefix_caching": True,
            "chunked_prefill": True,
            "attention_backend": "FLASH_ATTN",
        },
        "capture_contract": {
            "api": "LLM.chat",
            "enable_return_routed_experts": True,
            "route_shape": "[prompt_tokens + generated_tokens - 1, layers, topk]",
            "decode_policy": (
                f"fixed_override_{args.decode_override}"
                if args.decode_override is not None
                else "exact_trace_output_length"
            ),
            "contains_prompt_text": False,
            "fixture_sha256": sha256(args.fixture),
            "routes_npz": str(args.routes),
        },
        "requests": request_summaries,
        "prefix_pairs": prefix_pairs,
        "phases": {
            "prefill": phase_summary(prefill_routes),
            "decode": phase_summary(decode_routes),
        },
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    print(
        json.dumps(
            {
                "requests": len(rows),
                "prefill_tokens": payload["phases"]["prefill"]["token_count"],
                "decode_tokens": payload["phases"]["decode"]["token_count"],
                "prefix_pairs": prefix_pairs,
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()
