"""Device plumbing shared by the deep_ep tests.

`DEEP_EP_TEST_DEVICE=neuron` runs every rank on its own NeuronCore (core `DEEP_EP_TEST_CORE_BASE + rank`) and
hands the buffers a torch-neuronx ``neuron`` process group, exercising the real collectives. Inputs are generated
and references are computed on the **host** (gloo default group, CPU copies), then compared against the device
results copied back, so a device bug cannot cancel out between the code under test and its reference.
Default (`cpu`): everything on the host with gloo, as before.

FP8 tensors cross the host/device boundary as `uint8` views (indexing / cat on float8 is wrong on the neuron
device, see tests/neuron/device/test_device_primitives.py); strides are preserved (column-major SF inputs).
"""

import os

import torch
import torch.distributed as dist

DEVICE = os.environ.get('DEEP_EP_TEST_DEVICE', 'cpu')
_FP8 = (torch.float8_e4m3fn, torch.float8_e5m2)


def init_dist(local_rank: int, num_ranks: int, port: int):
    """Default gloo PG (host references) + the group handed to the buffers.

    Mirrors DeepInfer's own layout (`worker/base.py`: gloo default PG, `PLATFORM.distributed_backend` subgroups).
    torch-neuronx sets up the runtime when the first ``neuron`` group is created and asserts that no device op ran
    before, so NEURON_RT_VISIBLE_CORES is pinned here and nothing touches the device until the group exists.
    """
    os.environ.setdefault('MASTER_ADDR', '127.0.0.1')
    os.environ.setdefault('MASTER_PORT', str(port))
    if DEVICE == 'neuron':
        os.environ['NEURON_RT_VISIBLE_CORES'] = str(int(os.environ.get('DEEP_EP_TEST_CORE_BASE', '0')) + local_rank)
        # Same process model as DeepInfer (deepinfer/platform.py PlatformNeuron): one core per process, which
        # torch-neuronx then exposes as neuron:0 of a 1-device local world on every rank.
        os.environ['LOCAL_RANK'] = '0'
        os.environ['LOCAL_WORLD_SIZE'] = '1'
    dist.init_process_group('gloo', rank=local_rank, world_size=num_ranks)
    if DEVICE == 'neuron':
        import torch_neuronx  # noqa: F401  (registers the "neuron" backend)
        group = dist.new_group(list(range(num_ranks)), backend='neuron')
    else:
        group = dist.group.WORLD
    return group


def to_dev(t):
    """Host tensor -> DEVICE, same dtype/shape/strides (identity on cpu)."""
    if t is None or DEVICE == 'cpu':
        return t
    if isinstance(t, tuple):
        return tuple(to_dev(x) for x in t)
    is_fp8 = t.dtype in _FP8
    src = t.view(torch.uint8) if is_fp8 else t
    out = torch.empty_strided(src.size(), src.stride(), dtype=src.dtype, device=DEVICE)
    out.copy_(src)
    return out.view(t.dtype) if is_fp8 else out


def to_host(t):
    """DEVICE tensor -> host (identity on cpu); FP8 through a uint8 view."""
    if t is None or DEVICE == 'cpu':
        return t
    if isinstance(t, tuple):
        return tuple(to_host(x) for x in t)
    if t.dtype in _FP8:
        return t.view(torch.uint8).cpu().view(t.dtype)
    return t.cpu()