"""Benchmark CP shared-KV prefetch collective layouts.

This benchmark isolates the current CP shared-KV prefetch merge pattern:

  dense_all_reduce:
    each rank materializes its owner-lane pages into a full dense prefix
    buffer, zeros non-owned pages, then all-reduces the full dense buffer.

Against two all-gather candidates:

  packed_all_gather_scatter:
    each rank gathers only owner-lane pages, all-gathers owner-packed pages,
    then scatters them back into the old dense slot layout.

  packed_all_gather_rank_major:
    each rank gathers only owner-lane pages and keeps the all-gather output in
    rank-major owner-packed order.  This models the intended design where
    dense_locs/page_inverse are remapped instead of scattering back.

Run on the target CUDA machine, for example:

  torchrun --nproc_per_node=8 \\
    benchmark/hicache/bench_cp_shared_kv_prefetch_collective.py \\
    --payload both --tokens 4096 8192 16384 32768 40320 65536

Do not use this as an end-to-end throughput benchmark.  It is a microbenchmark
for materialize + collective + optional scatter cost.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import statistics
import time
from dataclasses import asdict, dataclass
from typing import Callable, Iterable

import torch
import torch.distributed as dist


@dataclass(frozen=True)
class CaseConfig:
    payload: str
    prefix_tokens: int
    page_size: int
    kv_dim: int
    index_page_bytes: int
    dtype: str
    warmup: int
    repeat: int


@dataclass(frozen=True)
class BenchResult:
    payload: str
    prefix_tokens: int
    prefix_pages: int
    path: str
    full_mib: float
    packed_mib_per_rank: float
    gpu_ms_avg: float
    gpu_ms_p50: float
    gpu_ms_min: float
    gpu_ms_max_rank_avg: float
    cpu_ms_avg: float
    cpu_ms_p50: float
    cpu_ms_max_rank_avg: float


def _rank_env() -> tuple[int, int, int]:
    rank = int(os.environ.get("RANK", "0"))
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    local_rank = int(os.environ.get("LOCAL_RANK", str(rank)))
    return rank, world_size, local_rank


def _dtype(name: str) -> torch.dtype:
    if name == "bf16":
        return torch.bfloat16
    if name == "fp16":
        return torch.float16
    if name == "fp32":
        return torch.float32
    raise ValueError(f"unsupported dtype: {name}")


def _payload_shape_and_dtype(cfg: CaseConfig) -> tuple[tuple[int, ...], torch.dtype]:
    if cfg.payload == "mla":
        return (cfg.page_size, cfg.kv_dim), _dtype(cfg.dtype)
    if cfg.payload == "index":
        return (cfg.index_page_bytes,), torch.uint8
    raise ValueError(f"unsupported payload: {cfg.payload}")


def _owner_counts(prefix_pages: int, world_size: int) -> list[int]:
    return [len(range(rank, prefix_pages, world_size)) for rank in range(world_size)]


def _make_source_pages(
    *,
    max_owned_pages: int,
    page_shape: tuple[int, ...],
    dtype: torch.dtype,
    device: torch.device,
    rank: int,
) -> torch.Tensor:
    # Page 0 is kept as a sentinel to match CP shared-KV physical page usage.
    shape = (max_owned_pages + 1, *page_shape)
    if dtype is torch.uint8:
        src = torch.empty(shape, dtype=dtype, device=device)
        src.random_(0, 127)
    else:
        src = torch.empty(shape, dtype=dtype, device=device)
        src.normal_(mean=float(rank + 1), std=0.01)
    src[0].zero_()
    return src


def _copy_owned_to_dense(
    *,
    dense_pages: torch.Tensor,
    src_pages: torch.Tensor,
    owned_slots: torch.Tensor,
    owned_count: int,
) -> None:
    if owned_count <= 0:
        return
    dense_pages[owned_slots].copy_(src_pages[1 : owned_count + 1])


def _copy_owned_to_packed(
    *,
    local_packed: torch.Tensor,
    src_pages: torch.Tensor,
    owned_count: int,
) -> None:
    local_packed.zero_()
    if owned_count > 0:
        local_packed[:owned_count].copy_(src_pages[1 : owned_count + 1])


def _scatter_rank_major_to_dense(
    *,
    dense_pages: torch.Tensor,
    gathered: torch.Tensor,
    prefix_pages: int,
    owner_counts: list[int],
    world_size: int,
    device: torch.device,
) -> None:
    dense_pages.zero_()
    for owner_rank, count in enumerate(owner_counts):
        if count <= 0:
            continue
        slots = torch.arange(
            owner_rank,
            prefix_pages,
            world_size,
            device=device,
            dtype=torch.long,
        )
        dense_pages[slots].copy_(gathered[owner_rank, :count])


def _time_path(
    fn: Callable[[], None],
    *,
    warmup: int,
    repeat: int,
    stream: torch.cuda.Stream,
    group: dist.ProcessGroup,
    device: torch.device,
) -> tuple[list[float], list[float]]:
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize(device)
    dist.barrier(group=group)

    gpu_ms: list[float] = []
    cpu_ms: list[float] = []
    for _ in range(repeat):
        dist.barrier(group=group)
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        torch.cuda.synchronize(device)
        t0 = time.perf_counter()
        with torch.cuda.stream(stream):
            start.record(stream)
            fn()
            end.record(stream)
        end.synchronize()
        t1 = time.perf_counter()
        gpu_ms.append(float(start.elapsed_time(end)))
        cpu_ms.append((t1 - t0) * 1000.0)
    dist.barrier(group=group)
    return gpu_ms, cpu_ms


def _summarize_across_ranks(
    *,
    local_gpu_ms: list[float],
    local_cpu_ms: list[float],
    group: dist.ProcessGroup,
    device: torch.device,
) -> tuple[dict[str, float], dict[str, float]]:
    local = torch.tensor(
        [
            statistics.mean(local_gpu_ms),
            statistics.median(local_gpu_ms),
            min(local_gpu_ms),
            statistics.mean(local_cpu_ms),
            statistics.median(local_cpu_ms),
        ],
        dtype=torch.float64,
        device=device,
    )
    world_size = dist.get_world_size(group=group)
    gathered = torch.empty((world_size, local.numel()), dtype=local.dtype, device=device)
    dist.all_gather_into_tensor(gathered, local, group=group)
    gathered_cpu = gathered.cpu()
    gpu = {
        "avg": float(gathered_cpu[:, 0].mean().item()),
        "p50": float(gathered_cpu[:, 1].mean().item()),
        "min": float(gathered_cpu[:, 2].mean().item()),
        "max_rank_avg": float(gathered_cpu[:, 0].max().item()),
    }
    cpu = {
        "avg": float(gathered_cpu[:, 3].mean().item()),
        "p50": float(gathered_cpu[:, 4].mean().item()),
        "max_rank_avg": float(gathered_cpu[:, 3].max().item()),
    }
    return gpu, cpu


def _bench_case(
    cfg: CaseConfig,
    *,
    group: dist.ProcessGroup,
    device: torch.device,
    rank: int,
    world_size: int,
) -> list[BenchResult]:
    prefix_pages = math.ceil(cfg.prefix_tokens / cfg.page_size)
    owner_counts = _owner_counts(prefix_pages, world_size)
    owned_count = owner_counts[rank]
    max_owned_pages = max(owner_counts)
    page_shape, dtype = _payload_shape_and_dtype(cfg)
    src_pages = _make_source_pages(
        max_owned_pages=max_owned_pages,
        page_shape=page_shape,
        dtype=dtype,
        device=device,
        rank=rank,
    )

    dense_pages = torch.empty((prefix_pages, *page_shape), dtype=dtype, device=device)
    local_packed = torch.empty(
        (max_owned_pages, *page_shape), dtype=dtype, device=device
    )
    gathered = torch.empty(
        (world_size, max_owned_pages, *page_shape), dtype=dtype, device=device
    )
    owned_slots = torch.arange(
        rank,
        prefix_pages,
        world_size,
        device=device,
        dtype=torch.long,
    )
    stream = torch.cuda.Stream(device=device)
    full_mib = dense_pages.numel() * dense_pages.element_size() / (1024 * 1024)
    packed_mib = local_packed.numel() * local_packed.element_size() / (1024 * 1024)

    def dense_all_reduce() -> None:
        dense_pages.zero_()
        _copy_owned_to_dense(
            dense_pages=dense_pages,
            src_pages=src_pages,
            owned_slots=owned_slots,
            owned_count=owned_count,
        )
        dist.all_reduce(dense_pages, op=dist.ReduceOp.SUM, group=group)

    def packed_all_gather_scatter() -> None:
        _copy_owned_to_packed(
            local_packed=local_packed,
            src_pages=src_pages,
            owned_count=owned_count,
        )
        dist.all_gather_into_tensor(gathered, local_packed, group=group)
        _scatter_rank_major_to_dense(
            dense_pages=dense_pages,
            gathered=gathered,
            prefix_pages=prefix_pages,
            owner_counts=owner_counts,
            world_size=world_size,
            device=device,
        )

    def packed_all_gather_rank_major() -> None:
        _copy_owned_to_packed(
            local_packed=local_packed,
            src_pages=src_pages,
            owned_count=owned_count,
        )
        dist.all_gather_into_tensor(gathered, local_packed, group=group)

    paths: list[tuple[str, Callable[[], None]]] = [
        ("dense_all_reduce", dense_all_reduce),
        ("packed_all_gather_scatter", packed_all_gather_scatter),
        ("packed_all_gather_rank_major", packed_all_gather_rank_major),
    ]
    results: list[BenchResult] = []
    for name, fn in paths:
        gpu_ms, cpu_ms = _time_path(
            fn,
            warmup=cfg.warmup,
            repeat=cfg.repeat,
            stream=stream,
            group=group,
            device=device,
        )
        gpu, cpu = _summarize_across_ranks(
            local_gpu_ms=gpu_ms,
            local_cpu_ms=cpu_ms,
            group=group,
            device=device,
        )
        results.append(
            BenchResult(
                payload=cfg.payload,
                prefix_tokens=cfg.prefix_tokens,
                prefix_pages=prefix_pages,
                path=name,
                full_mib=full_mib,
                packed_mib_per_rank=packed_mib,
                gpu_ms_avg=gpu["avg"],
                gpu_ms_p50=gpu["p50"],
                gpu_ms_min=gpu["min"],
                gpu_ms_max_rank_avg=gpu["max_rank_avg"],
                cpu_ms_avg=cpu["avg"],
                cpu_ms_p50=cpu["p50"],
                cpu_ms_max_rank_avg=cpu["max_rank_avg"],
            )
        )
    return results


def _print_result(result: BenchResult) -> None:
    print(json.dumps(asdict(result), sort_keys=True), flush=True)
    print(
        f"{result.payload:5s} tokens={result.prefix_tokens:6d} "
        f"pages={result.prefix_pages:5d} {result.path:28s} "
        f"gpu_avg={result.gpu_ms_avg:8.3f}ms "
        f"gpu_max_rank={result.gpu_ms_max_rank_avg:8.3f}ms "
        f"cpu_avg={result.cpu_ms_avg:8.3f}ms "
        f"full={result.full_mib:8.1f}MiB "
        f"packed/rank={result.packed_mib_per_rank:8.1f}MiB",
        flush=True,
    )


def _payloads(value: str) -> Iterable[str]:
    if value == "both":
        return ("mla", "index")
    return (value,)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Benchmark CP shared-KV prefetch all-reduce vs all-gather layouts."
    )
    parser.add_argument("--payload", choices=("mla", "index", "both"), default="both")
    parser.add_argument(
        "--tokens",
        type=int,
        nargs="+",
        default=[4096, 8192, 16384, 32768, 40320, 65536],
        help="Prefix token counts to benchmark.",
    )
    parser.add_argument("--page-size", type=int, default=64)
    parser.add_argument("--kv-dim", type=int, default=576)
    parser.add_argument(
        "--index-page-bytes",
        type=int,
        default=8448,
        help="Bytes per index page. GLM5 NSA observed default is 8448.",
    )
    parser.add_argument("--dtype", choices=("bf16", "fp16", "fp32"), default="bf16")
    parser.add_argument("--warmup", type=int, default=5)
    parser.add_argument("--repeat", type=int, default=20)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    rank, world_size, local_rank = _rank_env()
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for this benchmark.")
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    dist.init_process_group("nccl", init_method="env://")
    group = dist.group.WORLD

    if rank == 0:
        print(
            "CP shared-KV prefetch collective benchmark "
            f"world_size={world_size} tokens={args.tokens} payload={args.payload}",
            flush=True,
        )

    try:
        for payload in _payloads(args.payload):
            for tokens in args.tokens:
                cfg = CaseConfig(
                    payload=payload,
                    prefix_tokens=int(tokens),
                    page_size=args.page_size,
                    kv_dim=args.kv_dim,
                    index_page_bytes=args.index_page_bytes,
                    dtype=args.dtype,
                    warmup=args.warmup,
                    repeat=args.repeat,
                )
                for result in _bench_case(
                    cfg,
                    group=group,
                    device=device,
                    rank=rank,
                    world_size=world_size,
                ):
                    if rank == 0:
                        _print_result(result)
    finally:
        dist.barrier(group=group)
        dist.destroy_process_group()


if __name__ == "__main__":
    main()