# SPDX-License-Identifier: Apache-2.0
# Derived from vLLM; modified by DeepSeek.

from typing import Any, NamedTuple

import torch
import torch.distributed as dist

from deepinfer.model_executor.parallel_utils.parallel_state import PG, cp_broadcast_object


def gather(
    input_: torch.Tensor,
    group: dist.ProcessGroup,
    dst: int = 0,
    dim: int = -1,
) -> torch.Tensor | None:
    """Gather input_ from all ranks of group and concat along dim; returns None on non-dst ranks."""
    if group.size() == 1:
        return input_
    assert -input_.dim() <= dim < input_.dim(), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
    is_dst = group.rank() == dst
    gather_list = [torch.empty_like(input_) for _ in range(group.size())] if is_dst else None
    dist.gather(input_, gather_list, dst=dist.get_global_rank(group, dst), group=group)
    return torch.cat(gather_list, dim=dim) if is_dst else None


def cp_broadcast(input_: torch.Tensor, src: int = 0) -> torch.Tensor:
    """Broadcast input_ in place across the CP group."""
    cp_group = PG.context_parallel_group()
    if cp_group.size() > 1:
        dist.broadcast(input_, src=dist.get_global_rank(cp_group, src), group=cp_group)
    return input_


def cp_broadcast_multi_tensor(inputs: list[torch.Tensor], src: int = 0) -> None:
    """Broadcast multiple tensors in place across the CP group."""
    cp_group = PG.context_parallel_group()
    if cp_group.size() == 1:
        return
    # TODO: Maybe use deepep multi_tensor_broadcast later.
    for tensor in inputs:
        dist.broadcast(tensor, src=dist.get_global_rank(cp_group, src), group=cp_group)


class TensorMetadata(NamedTuple):
    dtype: torch.dtype
    size: torch.Size


def broadcast_tensor_dict(
    tensor_dict: dict[Any, Any] | None = None,
    src: int = 0,
) -> dict[Any, Any]:
    """Broadcast a dict across the CP group: cuda tensors via NCCL, everything else via shm object broadcast."""
    cp_group = PG.context_parallel_group()
    if cp_group.size() == 1:
        return tensor_dict
    assert src == 0, f"Invalid src rank ({src})"

    tensors = []
    if cp_group.rank() == src:
        assert isinstance(tensor_dict, dict), f"Expecting a dictionary, got {type(tensor_dict)}"
        metadata_list = []
        for key, value in tensor_dict.items():
            if isinstance(value, torch.Tensor):
                assert value.is_cuda, f"Tensor {key}: {value} is not on cuda. Currently we only support broadcasting tensors on cuda."
                metadata_list.append((key, TensorMetadata(value.dtype, value.size())))
                tensors.append(value)
            else:
                metadata_list.append((key, value))
        cp_broadcast_object(metadata_list)
    else:
        tensor_dict = {}
        for key, value in cp_broadcast_object(None):
            if isinstance(value, TensorMetadata):
                tensor_dict[key] = torch.empty(value.size, dtype=value.dtype, device="cuda")
                tensors.append(tensor_dict[key])
            else:
                tensor_dict[key] = value
    cp_broadcast_multi_tensor(tensors, src)
    return tensor_dict