import torch
from contextlib import contextmanager, nullcontext
from torch.distributed import Backend
from typing import Any, Dict, List, Optional, Tuple, Union
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.distributed.parallel_state import (GroupCoordinator,
                                             GraphCaptureContext)
from vllm.logger import init_logger

logger = init_logger(__name__)


vllm__distributed__parallel_state__GroupCoordinator____init____org = GroupCoordinator.__init__


def vllm__distributed__parallel_state__GroupCoordinator____init__(
    self,
    group_ranks: List[List[int]],
    local_rank: int,
    torch_distributed_backend: Union[str, Backend],
    use_pynccl: bool,
    use_custom_allreduce: bool,
    use_tpu_communicator: bool,
    use_hpu_communicator: bool,
    use_xpu_communicator: bool,
    use_message_queue_broadcaster: bool = False,
    group_name: Optional[str] = None,
):
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: disable pynccl and custom_allreduce by default
    '''
    if use_pynccl or use_custom_allreduce:
        logger.debug(f"Disable pynccl and custom_allreduce when using MLU backend.")

    vllm__distributed__parallel_state__GroupCoordinator____init____org(
        self=self,
        group_ranks=group_ranks,
        local_rank=local_rank,
        torch_distributed_backend=torch_distributed_backend,
        use_pynccl=False,
        use_custom_allreduce=False,
        use_tpu_communicator=use_tpu_communicator,
        use_hpu_communicator=use_hpu_communicator,
        use_xpu_communicator=use_xpu_communicator,
        use_message_queue_broadcaster=use_message_queue_broadcaster,
        group_name=group_name
    )
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


def vllm__distributed__parallel_state__GroupCoordinator__gather(
        self,
        input_: torch.Tensor,
        dst: int = 0,
        dim: int = -1) -> Optional[torch.Tensor]:
    """
    NOTE: We assume that the input tensor is on the same device across
    all the ranks.
    NOTE: `dst` is the local rank of the destination rank.
    """
    world_size = self.world_size
    # Bypass the function if we are using only 1 GPU.
    if world_size == 1:
        return input_
    assert -input_.dim() <= dim < input_.dim(), (
        f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")
    if dim < 0:
        # Convert negative dim to positive.
        dim += input_.dim()
    if self.xpu_communicator is not None and \
            not self.xpu_communicator.disabled:
        return self.xpu_communicator.gather(input_, self.rank_in_group,
                                            dst, dim)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: use p2p communication to reduce gather host time for non driver worker.
    NOTE: this hijack function should be REMOVED when torch upgrade to 2.4.0
    '''
    rank = self.rank_in_group

    gather_list = None
    if rank == dst:
        gather_list = [
            torch.empty_like(input_) for _ in range(dst)
        ] + [input_] + [
            torch.empty_like(input_) for _ in range(dst + 1, world_size)
        ]

    send_recv_op_list = []
    if rank != dst:
        op = torch.distributed.P2POp(torch.distributed.isend,
                                     input_,
                                     self.ranks[dst],
                                     group=self.device_group)
        send_recv_op_list.append(op)
    else:
        for r in range(0, world_size):
            if r == dst:
                continue
            op = torch.distributed.P2POp(torch.distributed.irecv,
                                         gather_list[r],
                                         self.ranks[r],
                                         group=self.device_group)
            send_recv_op_list.append(op)
    reqs = torch.distributed.batch_isend_irecv(send_recv_op_list)
    for req in reqs:
        req.wait()
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    if self.rank_in_group == dst:
        output_tensor = torch.cat(gather_list, dim=dim)
    else:
        output_tensor = None
    return output_tensor


MluHijackObject.apply_hijack(GroupCoordinator,
                             GroupCoordinator.__init__,
                             vllm__distributed__parallel_state__GroupCoordinator____init__)
MluHijackObject.apply_hijack(GroupCoordinator,
                             GroupCoordinator.gather,
                             vllm__distributed__parallel_state__GroupCoordinator__gather)