#!/usr/bin/env python3
"""8-rank GPU byte-exactness test: compose legacy vs v2 (Step A).

Runs the REAL materialize_prefix_and_reuse_current_{kv,index}_page_slots on
every rank with real NCCL collectives and (for v2) real tai-kernel CUDA-IPC
gathers, and asserts the composed dense buffers and locs are byte-identical
between the legacy per-span path and compose_v2.

Run inside the g0034 cjy-glm5-new container:
  cd /mnt/beegfs/syh/sglang-stable && \
  SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 \
  PYTHONPATH=python:/mnt/beegfs/syh/tai-kernel/python \
  torchrun --nproc-per-node=8 test/manual/test_cp_shared_kv_compose_v2_8rank.py
"""
from __future__ import annotations

import os

import torch
import torch.distributed as dist

from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout


class _CpGroupShim:
    """Minimal stand-in for the attention CP GroupCoordinator."""

    def __init__(self, group: dist.ProcessGroup) -> None:
        self.device_group = group
        self.world_size = dist.get_world_size(group)
        self.unique_name = "compose_v2_test"
        self.pynccl_comm = None

    def all_gather_into_tensor(self, out: torch.Tensor, t: torch.Tensor) -> None:
        dist.all_gather_into_tensor(out, t, group=self.device_group)

    def all_reduce(self, t: torch.Tensor) -> torch.Tensor:
        dist.all_reduce(t, group=self.device_group)
        return t


def _build_scenario(rank: int, cp_size: int, device: torch.device):
    """bs=4 with mixed prefix/extend lengths, extends crossing page bounds."""
    page_size = 64
    batch_size = 4
    prefix_lens = [640, 1280, 320, 640]
    extend_lens = [95, 130, 64, 200]
    kv_dim = 656  # fp8 layout bytes/token

    from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
        build_in_seq_page_compute_owners,
    )

    g = torch.Generator().manual_seed(20260612)
    logical_rows = []
    current_locs_all = []
    next_page = 1
    prefix_pages_by_req = []
    for req_id, (prefix_len, extend_len) in enumerate(zip(prefix_lens, extend_lens)):
        prefix_pages = prefix_len // page_size
        prefix_pages_by_req.append(prefix_pages)
        req_pages = list(range(next_page, next_page + prefix_pages))
        next_page += prefix_pages
        owners = build_in_seq_page_compute_owners(
            extend_len=extend_len,
            extend_prefix_len=prefix_len,
            page_size=page_size,
            cp_size=cp_size,
        )
        current_pages = []
        for owner in owners:
            # pick a logical page owned by `owner`: owner = (page-1) % cp_size
            candidate = next_page
            while (candidate - 1) % cp_size != int(owner):
                candidate += 1
            current_pages.append(candidate)
            next_page = candidate + 1
        req_pages.extend(current_pages)
        remaining = extend_len
        for page_offset, logical_page in enumerate(current_pages):
            valid = min(page_size, remaining)
            for off in range(valid):
                current_locs_all.append(
                    (req_id, logical_page * page_size + off)
                )
            remaining -= valid
        logical_rows.append(req_pages)

    max_pages = max(len(r) for r in logical_rows)
    logical_pages = torch.zeros((batch_size, max_pages), dtype=torch.int64)
    for req_id, pages in enumerate(logical_rows):
        logical_pages[req_id, : len(pages)] = torch.tensor(pages, dtype=torch.int64)
    logical_pages = logical_pages.to(device)

    # logical locs per request row (token granularity, -1 padded)
    max_tokens = max(
        p + e for p, e in zip(prefix_lens, extend_lens)
    )
    logical_locs = torch.full((batch_size, max_tokens), -1, dtype=torch.int64)
    loc_req_id_rows = []
    for req_id, pages in enumerate(logical_rows):
        total = prefix_lens[req_id] + extend_lens[req_id]
        locs = []
        for token_idx in range(total):
            page = pages[token_idx // page_size]
            locs.append(page * page_size + token_idx % page_size)
        logical_locs[req_id, : len(locs)] = torch.tensor(locs, dtype=torch.int64)
    logical_locs = logical_locs.to(device)

    physical_pages = (next_page // cp_size + 3)
    kv_cache = torch.zeros(
        (physical_pages * page_size, 1, kv_dim), dtype=torch.uint8, device=device
    )
    layout = CpSharedKVLayout(page_size=page_size, cp_size=cp_size, cp_rank=rank)

    # Fill this rank's owned PREFIX pages with deterministic payloads.
    for req_id, pages in enumerate(logical_rows):
        for slot, logical_page in enumerate(pages[: prefix_pages_by_req[req_id]]):
            if (logical_page - 1) % cp_size != rank:
                continue
            phys_page = (logical_page - 1) // cp_size + 1
            payload = (
                torch.arange(page_size * kv_dim, dtype=torch.int64)
                + logical_page * 131
            ).remainder_(251).to(torch.uint8)
            kv_cache[
                phys_page * page_size : (phys_page + 1) * page_size
            ] = payload.view(page_size, 1, kv_dim).to(device)

    # Current rows owned by this rank (writer = page owner in this scenario,
    # so the symm writer list is the storage owner per current page).
    current_page_writer_ranks = []
    for req_pages, prefix_pages in zip(logical_rows, prefix_pages_by_req):
        for page in req_pages[prefix_pages:]:
            current_page_writer_ranks.append((page - 1) % cp_size)
    cur_locs = [
        loc for req_id, loc in current_locs_all
        if (loc // page_size - 1) % cp_size == rank
    ]
    cur_req = [
        req_id for req_id, loc in current_locs_all
        if (loc // page_size - 1) % cp_size == rank
    ]
    current_locs = torch.tensor(cur_locs, dtype=torch.int64, device=device)
    current_req_id = torch.tensor(cur_req, dtype=torch.int64, device=device)
    current_kv = (
        (
            torch.arange(len(cur_locs), dtype=torch.int64).view(-1, 1, 1) * 7
            + torch.arange(kv_dim, dtype=torch.int64).view(1, 1, -1)
            + rank * 31
        )
        .remainder_(249)
        .to(torch.uint8)
        .to(device)
    )

    prefix_slot_spans = runtime.build_batch_prefix_slot_spans(
        logical_pages=logical_pages,
        prefix_lens_cpu=prefix_lens,
        page_size=page_size,
    )
    current_slot_spans = runtime.build_batch_current_slot_spans(
        logical_pages=logical_pages,
        prefix_lens_cpu=prefix_lens,
        extend_lens_cpu=extend_lens,
        page_size=page_size,
    )
    slot_remap = runtime.build_shared_token_kv_slot_remap(
        kv_cache,
        logical_locs,
        logical_pages,
        layout,
        page_size,
    )
    loc_req_id = torch.repeat_interleave(
        torch.arange(batch_size, dtype=torch.int64, device=device),
        torch.tensor(
            [max_tokens] * batch_size, dtype=torch.int64, device=device
        ),
    )

    return dict(
        kv_cache=kv_cache,
        logical_locs=logical_locs.reshape(-1),
        loc_req_id=loc_req_id,
        current_kv=current_kv,
        current_locs=current_locs,
        current_req_id=current_req_id,
        slot_remap=slot_remap,
        layout=layout,
        page_size=page_size,
        prefix_slot_spans=prefix_slot_spans,
        current_slot_spans=current_slot_spans,
        current_page_writer_ranks=current_page_writer_ranks,
    )


def _compose(s, layer_id: int, *, writers: list[int] | None = None):
    return runtime.materialize_prefix_and_reuse_current_kv_page_slots(
        kv_cache=s["kv_cache"],
        logical_locs=s["logical_locs"],
        current_kv_cache=s["current_kv"],
        current_locs=s["current_locs"],
        slot_remap=s["slot_remap"],
        layout=s["layout"],
        page_size=s["page_size"],
        prefix_pages=0,
        loc_req_id=s["loc_req_id"],
        current_req_id=s["current_req_id"],
        prefix_slot_spans=s["prefix_slot_spans"],
        current_slot_spans=s["current_slot_spans"],
        layer_id=layer_id,
        current_page_writer_ranks=writers,
    )


def main() -> None:
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    world = dist.get_world_size()
    local_rank = int(os.environ.get("LOCAL_RANK", rank))
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)

    shim = _CpGroupShim(dist.group.WORLD)
    runtime.get_attention_cp_group = lambda: shim  # type: ignore[assignment]

    s = _build_scenario(rank, world, device)

    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False):
        ref_kv, ref_locs = _compose(s, layer_id=0)
    torch.cuda.synchronize()
    dist.barrier()

    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(True):
        v2_kv, v2_locs = _compose(s, layer_id=0)
    torch.cuda.synchronize()
    dist.barrier()

    assert torch.equal(ref_locs, v2_locs), (
        f"rank{rank}: locs mismatch legacy vs v2"
    )
    if not torch.equal(ref_kv, v2_kv):
        diff = (ref_kv != v2_kv).any(dim=-1).any(dim=-1)
        bad_rows = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist()
        raise AssertionError(
            f"rank{rank}: dense kv mismatch at rows {bad_rows} "
            f"(of {int(diff.sum())} differing rows)"
        )

    # Cross-rank: every rank must hold the SAME composed buffer.
    ref_sum = ref_kv.to(torch.float64).sum()
    sums = torch.zeros(world, dtype=torch.float64, device=device)
    dist.all_gather_into_tensor(
        sums, ref_sum.reshape(1).to(device)
    )
    assert torch.allclose(sums, sums[0].expand_as(sums)), (
        f"rank{rank}: composed buffers differ across ranks: {sums.tolist()}"
    )

    if rank == 0:
        print(
            "PASS: legacy vs compose_v2 byte-identical "
            f"(dense rows={int(ref_kv.shape[0])}, world={world})",
            flush=True,
        )

    # ---- Step B: symm exchange (arena + barrier + peer gather, zero NCCL
    # in the current-page phase). Multiple layers exercise parity halves. ----
    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(
        True
    ), envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(
        True
    ), envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(True):
        for layer_id in range(4):
            symm_kv, symm_locs = _compose(
                s, layer_id=layer_id, writers=s["current_page_writer_ranks"]
            )
            torch.cuda.synchronize()
            assert torch.equal(ref_locs, symm_locs), (
                f"rank{rank} layer{layer_id}: symm locs mismatch"
            )
            if not torch.equal(ref_kv, symm_kv):
                diff = (ref_kv != symm_kv).any(dim=-1).any(dim=-1)
                bad = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist()
                raise AssertionError(
                    f"rank{rank} layer{layer_id}: symm dense kv mismatch at "
                    f"rows {bad} (of {int(diff.sum())})"
                )
    dist.barrier()
    from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
        get_compose_arena,
    )

    assert get_compose_arena(device).symm_ready, "symm slab was not registered"
    if rank == 0:
        print(
            "PASS: symm compose byte-identical to v2 across 4 layers "
            "(arena registered, barrier + peer gather engaged)",
            flush=True,
        )
    dist.destroy_process_group()


if __name__ == "__main__":
    main()