# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for SeqAllToAll4D, SeqAllToAll5D, and RingComm communication primitives.

CPU tests use gloo on CPU tensors (no GPU). Nightly parity tests run the same
checks on real multi-GPU NCCL collectives.
"""

from __future__ import annotations

import os
from typing import Literal

import pytest
import torch

from tests.helpers.distributed import configure_cpu_collective_worker, start_collective_workers
from tests.helpers.mark import hardware_marks
from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D
from vllm_omni.diffusion.distributed.parallel_state import (
    destroy_distributed_env,
    get_sp_group,
    init_distributed_environment,
    initialize_model_parallel,
)
from vllm_omni.platforms import current_omni_platform

_L4_TWO_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=2)
_L4_FOUR_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=4)

DeviceKind = Literal["cpu", "cuda"]


def _update_environment_variables(envs_dict: dict[str, str]) -> None:
    for key, value in envs_dict.items():
        os.environ[key] = value


def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device:
    if device_kind == "cpu":
        return torch.device("cpu")
    return torch.device(f"{current_omni_platform.device_type}:{local_rank}")


def _init_worker(
    local_rank: int,
    world_size: int,
    master_port: int,
    device_kind: DeviceKind,
) -> None:
    _update_environment_variables(
        {
            "RANK": str(local_rank),
            "LOCAL_RANK": str(local_rank),
            "WORLD_SIZE": str(world_size),
            "MASTER_ADDR": "localhost",
            "MASTER_PORT": str(master_port),
        }
    )
    if device_kind == "cpu":
        configure_cpu_collective_worker()
    backend = "gloo" if device_kind == "cpu" else None
    init_distributed_environment(backend=backend)


def _close_tolerance(dtype: torch.dtype) -> tuple[float, float]:
    if dtype == torch.bfloat16:
        return 1e-3, 1e-3
    if dtype == torch.float16:
        return 1e-3, 1e-3
    return 1e-5, 1e-5


def _run_4d_identity(
    local_rank: int,
    world_size: int,
    dtype: torch.dtype,
    batch_size: int,
    seq_len_per_rank: int,
    num_heads: int,
    head_size: int,
    use_sync: bool,
    device_kind: DeviceKind,
    master_port: int,
) -> None:
    device = _worker_device(local_rank, device_kind)
    if device_kind == "cuda":
        current_omni_platform.set_device(device)

    _init_worker(local_rank, world_size, master_port, device_kind)
    initialize_model_parallel(ulysses_degree=world_size)
    sp_group = get_sp_group().ulysses_group

    torch.manual_seed(42 + local_rank)
    input_tensor = torch.randn(
        batch_size,
        seq_len_per_rank,
        num_heads,
        head_size,
        dtype=dtype,
        device=device,
    )
    original_input = input_tensor.clone()

    intermediate = SeqAllToAll4D.apply(sp_group, input_tensor, 2, 1, use_sync)
    expected_intermediate_shape = (
        batch_size,
        seq_len_per_rank * world_size,
        num_heads // world_size,
        head_size,
    )
    assert intermediate.shape == expected_intermediate_shape, (
        f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}"
    )

    output = SeqAllToAll4D.apply(sp_group, intermediate, 1, 2, use_sync)
    assert output.shape == original_input.shape, (
        f"Output shape mismatch: expected {original_input.shape}, got {output.shape}"
    )

    rtol, atol = _close_tolerance(dtype)
    torch.testing.assert_close(
        output,
        original_input,
        rtol=rtol,
        atol=atol,
        msg="Output does not match original input after two all-to-all operations",
    )
    destroy_distributed_env()


def _run_5d_identity(
    local_rank: int,
    world_size: int,
    dtype: torch.dtype,
    batch_size: int,
    seq_len_per_rank: int,
    num_heads: int,
    head_size: int,
    use_sync: bool,
    device_kind: DeviceKind,
    master_port: int,
) -> None:
    device = _worker_device(local_rank, device_kind)
    if device_kind == "cuda":
        current_omni_platform.set_device(device)

    _init_worker(local_rank, world_size, master_port, device_kind)
    initialize_model_parallel(ulysses_degree=world_size)
    sp_group = get_sp_group().ulysses_group

    torch.manual_seed(42 + local_rank)
    input_tensor = torch.randn(
        batch_size,
        seq_len_per_rank,
        3,
        num_heads,
        head_size,
        dtype=dtype,
        device=device,
    )
    original_input = input_tensor.clone()

    intermediate = SeqAllToAll5D.apply(sp_group, input_tensor, 3, 1, use_sync)
    expected_intermediate_shape = (
        batch_size,
        seq_len_per_rank * world_size,
        3,
        num_heads // world_size,
        head_size,
    )
    assert intermediate.shape == expected_intermediate_shape, (
        f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}"
    )

    output = SeqAllToAll5D.apply(sp_group, intermediate, 1, 3, use_sync)
    assert output.shape == original_input.shape, (
        f"Output shape mismatch: expected {original_input.shape}, got {output.shape}"
    )

    rtol, atol = _close_tolerance(dtype)
    torch.testing.assert_close(
        output,
        original_input,
        rtol=rtol,
        atol=atol,
        msg="Output does not match original input after two all-to-all operations",
    )
    destroy_distributed_env()


def _run_ring_p2p(
    local_rank: int,
    world_size: int,
    dtype: torch.dtype,
    batch_size: int,
    num_heads: int,
    head_size: int,
    device_kind: DeviceKind,
    master_port: int,
) -> None:
    device = _worker_device(local_rank, device_kind)
    if device_kind == "cuda":
        current_omni_platform.set_device(device)

    try:
        _init_worker(local_rank, world_size, master_port, device_kind)
        initialize_model_parallel(ring_degree=world_size)
        sp_group = get_sp_group()
        comm = RingComm(sp_group.ring_group)

        input_tensor = torch.full(
            (batch_size, num_heads, head_size),
            fill_value=float(local_rank + 1),
            dtype=dtype,
            device=device,
        )
        recv_tensor = comm.send_recv(input_tensor)
        comm.commit()
        comm.wait()

        prev_rank = (local_rank - 1 + world_size) % world_size
        expected_value = float(prev_rank + 1)
        expected_tensor = torch.full_like(recv_tensor, fill_value=expected_value)
        rtol, atol = _close_tolerance(dtype)
        torch.testing.assert_close(
            recv_tensor,
            expected_tensor,
            rtol=rtol,
            atol=atol,
            msg=f"[Rank {local_rank}] Ring P2P data mismatch",
        )
        if device_kind == "cpu":
            # P2P completion only synchronizes neighbors. Keep fast CPU ranks
            # alive until every rank has finished constructing and using groups.
            torch.distributed.barrier()
    finally:
        destroy_distributed_env()


def _spawn_4d_identity(
    *,
    world_size: int,
    dtype: torch.dtype,
    batch_size: int,
    seq_len_per_rank: int,
    num_heads: int,
    head_size: int,
    use_sync: bool,
    device_kind: DeviceKind,
    master_port: int,
) -> None:
    start_collective_workers(
        _run_4d_identity,
        args=(
            world_size,
            dtype,
            batch_size,
            seq_len_per_rank,
            num_heads,
            head_size,
            use_sync,
            device_kind,
            master_port,
        ),
        nprocs=world_size,
        device_kind=device_kind,
    )


def _spawn_5d_identity(
    *,
    world_size: int,
    dtype: torch.dtype,
    batch_size: int,
    seq_len_per_rank: int,
    num_heads: int,
    head_size: int,
    use_sync: bool,
    device_kind: DeviceKind,
    master_port: int,
) -> None:
    start_collective_workers(
        _run_5d_identity,
        args=(
            world_size,
            dtype,
            batch_size,
            seq_len_per_rank,
            num_heads,
            head_size,
            use_sync,
            device_kind,
            master_port,
        ),
        nprocs=world_size,
        device_kind=device_kind,
    )


def _spawn_ring_p2p(
    *,
    world_size: int,
    dtype: torch.dtype,
    batch_size: int,
    num_heads: int,
    head_size: int,
    device_kind: DeviceKind,
    master_port: int,
) -> None:
    start_collective_workers(
        _run_ring_p2p,
        args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port),
        nprocs=world_size,
        device_kind=device_kind,
    )


def _require_heads_divisible(num_heads: int, world_size: int) -> None:
    if num_heads % world_size != 0:
        pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")


def _require_gpus(world_size: int) -> None:
    available_gpus = current_omni_platform.get_device_count()
    if available_gpus < world_size:
        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")


# ---------------------------------------------------------------------------
# CPU: gloo collectives on CPU tensors (no GPU)
# ---------------------------------------------------------------------------


@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@pytest.mark.parametrize("world_size", [2, 4])
@pytest.mark.parametrize("use_sync", [False, True])
def test_4d_identity(world_size: int, use_sync: bool):
    _require_heads_divisible(8, world_size)
    _spawn_4d_identity(
        world_size=world_size,
        dtype=torch.float32,
        batch_size=2,
        seq_len_per_rank=8,
        num_heads=8,
        head_size=32,
        use_sync=use_sync,
        device_kind="cpu",
        master_port=29610,
    )


@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@pytest.mark.parametrize("world_size", [2, 4])
@pytest.mark.parametrize("use_sync", [False, True])
def test_5d_identity(world_size: int, use_sync: bool):
    _require_heads_divisible(8, world_size)
    _spawn_5d_identity(
        world_size=world_size,
        dtype=torch.float32,
        batch_size=2,
        seq_len_per_rank=8,
        num_heads=8,
        head_size=32,
        use_sync=use_sync,
        device_kind="cpu",
        master_port=29611,
    )


@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@pytest.mark.parametrize("world_size", [2, 4])
def test_ring_p2p(world_size: int):
    _spawn_ring_p2p(
        world_size=world_size,
        dtype=torch.float32,
        batch_size=2,
        num_heads=8,
        head_size=128,
        device_kind="cpu",
        master_port=29612,
    )


# ---------------------------------------------------------------------------
# Nightly: same checks on real multi-GPU NCCL collectives
# ---------------------------------------------------------------------------


@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.parametrize(
    "world_size",
    [
        pytest.param(2, marks=_L4_TWO_GPU),
        pytest.param(4, marks=_L4_FOUR_GPU),
    ],
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("use_sync", [False, True])
def test_4d_identity_parity(world_size: int, dtype: torch.dtype, use_sync: bool):
    _require_gpus(world_size)
    _require_heads_divisible(8, world_size)
    _spawn_4d_identity(
        world_size=world_size,
        dtype=dtype,
        batch_size=2,
        seq_len_per_rank=8,
        num_heads=8,
        head_size=32,
        use_sync=use_sync,
        device_kind="cuda",
        master_port=29500,
    )


@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.parametrize(
    "world_size",
    [
        pytest.param(2, marks=_L4_TWO_GPU),
        pytest.param(4, marks=_L4_FOUR_GPU),
    ],
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("use_sync", [False, True])
def test_5d_identity_parity(world_size: int, dtype: torch.dtype, use_sync: bool):
    _require_gpus(world_size)
    _require_heads_divisible(8, world_size)
    _spawn_5d_identity(
        world_size=world_size,
        dtype=dtype,
        batch_size=2,
        seq_len_per_rank=8,
        num_heads=8,
        head_size=32,
        use_sync=use_sync,
        device_kind="cuda",
        master_port=29502,
    )


@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.parametrize(
    "world_size",
    [
        pytest.param(2, marks=_L4_TWO_GPU),
        pytest.param(4, marks=_L4_FOUR_GPU),
    ],
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_ring_p2p_parity(world_size: int, dtype: torch.dtype):
    _require_gpus(world_size)
    _spawn_ring_p2p(
        world_size=world_size,
        dtype=dtype,
        batch_size=2,
        num_heads=8,
        head_size=128,
        device_kind="cuda",
        master_port=29501,
    )


def _run_fused_qkv_matches_three_4d(
    local_rank: int,
    world_size: int,
    device_kind: DeviceKind,
    master_port: int,
) -> None:
    """LingBot's fused (B, S/N, 3, H, D) exchange must equal three per-tensor 4D exchanges."""
    from vllm_omni.diffusion.distributed.comm import all_to_all_5D

    device = _worker_device(local_rank, device_kind)
    if device_kind == "cuda":
        current_omni_platform.set_device(device)
    _init_worker(local_rank, world_size, master_port, device_kind)
    initialize_model_parallel(ulysses_degree=world_size)
    sp_group = get_sp_group().ulysses_group
    try:
        batch, seq_per_rank, heads, head_size = 1, 6, 4 * world_size, 8
        # Distinguishable values: rank, q/k/v index, position, head, feature.
        base = torch.arange(seq_per_rank * heads * head_size, dtype=torch.float32, device=device).reshape(
            batch, seq_per_rank, heads, head_size
        )
        tensors = [base + 1000.0 * local_rank + 100000.0 * which for which in range(3)]
        separate = [SeqAllToAll4D.apply(sp_group, t, 2, 1, False) for t in tensors]
        fused = all_to_all_5D(torch.stack(tensors, dim=2), scatter_idx=3, gather_idx=1, group=sp_group)
        assert fused.shape == (batch, seq_per_rank * world_size, 3, heads // world_size, head_size)
        for which, expected in enumerate(separate):
            assert torch.equal(fused.unbind(2)[which], expected), f"q/k/v index {which} differs on rank {local_rank}"
    finally:
        destroy_distributed_env()


@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@pytest.mark.parametrize("world_size", [2, 4])
def test_fused_qkv_all_to_all_matches_three_4d_exchanges_cpu(world_size: int) -> None:
    start_collective_workers(
        _run_fused_qkv_matches_three_4d,
        args=(world_size, "cpu", 29613),
        nprocs=world_size,
        device_kind="cpu",
    )