"""BucketBuffer / BucketSession checks (subset of DeepEP-Ascend tests/bucket/*; CPU + gloo, or neuron, see dev.py).

Covers what `model_executor/parallel_utils/cp.py` uses: a session, `session.allocate`,
in-place and out-of-place `all_gather`, and the NotImplementedError stubs.
"""

import argparse
import os
import sys

import torch
import torch.distributed as dist
import torch.multiprocessing as mp

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor"))
import deep_ep  # noqa: E402
from dev import DEVICE, init_dist, to_dev, to_host  # noqa: E402


def worker(local_rank: int, num_ranks: int, args):
    group = init_dist(local_rank, num_ranks, args.port)
    rank = group.rank()

    alignment = deep_ep.get_num_allocation_alignment()
    assert alignment == 2 << 20 and deep_ep.get_num_rdma_alignment() == 512

    # Plan-based construction (upstream BufferAllocator path).
    plan = deep_ep.BufferAllocator()
    gathered = plan.allocate((num_ranks, 2, 8), torch.float32)
    buffer = deep_ep.BucketBuffer(group, plan, explicitly_destroy=True, device=torch.device(DEVICE))
    assert buffer.num_bytes == plan.num_bytes

    src_host = torch.arange(16, dtype=torch.float32).view(2, 8) + rank * 100
    src = to_dev(src_host)
    ref = torch.stack([src_host - rank * 100 + r * 100 for r in range(num_ranks)])  # host reference

    # Out-of-place: registered destination.
    gathered.fill_(float('nan'))
    result = buffer.all_gather(src, dsts=gathered, group=group).wait()
    assert result.ndim == 1 and torch.equal(to_host(result), ref.reshape(-1))

    # In-place: source is this rank's shard of a registered region.
    gathered.fill_(float('nan'))
    gathered[rank].copy_(src)
    result = buffer.all_gather(gathered[rank], group=group).wait()
    assert torch.equal(to_host(result), ref.reshape(-1))
    buffer.barrier(group=group)
    buffer.destroy()

    # Session path (what CPManager uses).
    buffer = deep_ep.BucketBuffer(group, 4 * alignment, explicitly_destroy=True,
                                 device=torch.device(DEVICE))
    with buffer.session() as session:
        region = session.allocate((num_ranks, 2, 8), torch.float32, group=group)
        shard = region[rank]
        shard.copy_(src)
        handle = buffer.all_gather(shard)
        out = handle.wait()
        assert torch.equal(to_host(out).view(num_ranks, 2, 8).flatten(0, 1), ref.reshape(-1, 8))

        # Unregistered source is staged into the session storage.
        unregistered = (src + 1000).contiguous()
        out = buffer.all_gather(unregistered).wait()
        ref2 = torch.stack([src_host - rank * 100 + r * 100 + 1000 for r in range(num_ranks)])
        assert torch.equal(to_host(out), ref2.reshape(-1))

        # Batched
        outs = buffer.all_gather([shard, shard]).wait()
        assert len(outs) == 2 and all(torch.equal(to_host(o), ref.reshape(-1)) for o in outs)

        for name in ('reduce_scatter', 'all_reduce'):
            try:
                getattr(buffer, name)(shard)
            except NotImplementedError:
                pass
            else:
                raise AssertionError(f'{name} should raise NotImplementedError')
    assert 'all_gather' not in buffer.__dict__
    buffer.destroy()

    if rank == 0:
        print(f'Bucket tests passed (world_size={num_ranks})', flush=True)
    dist.destroy_process_group()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--num-processes', type=int, default=2)
    parser.add_argument('--port', type=int, default=29551)
    args = parser.parse_args()
    if args.num_processes == 1:
        worker(0, 1, args)
    else:
        mp.spawn(worker, args=(args.num_processes, args), nprocs=args.num_processes)