import dataclasses
from typing import Any, Dict, Optional, Tuple, Union

import torch

from vllm.config import ObservabilityConfig, VllmConfig
from vllm.distributed.parallel_state import get_world_group
from vllm.sequence import ExecuteModelRequest
from vllm.worker.model_runner_base import (BroadcastableModelInput,
                                           ModelRunnerInputBase)
from vllm.worker.worker_base import (extract_previous_hidden_states,
                                     LocalOrDistributedWorkerBase,
                                     WorkerInput)
from vllm_mlu.mlu_hijack_utils import MluHijackObject


def broadcast_tensor_dict(
    tensor_dict: Optional[Dict[Any, Union[torch.Tensor, Any]]] = None,
    src: int = 0
):
    if not torch.distributed.is_initialized():
        return tensor_dict
    return get_world_group().broadcast_tensor_dict(tensor_dict, src)

def vllm__worker__worker_base__LocalOrDistributedWorkerBase___get_driver_input_and_broadcast(
    self, execute_model_req: ExecuteModelRequest
) -> Tuple[BroadcastableModelInput, WorkerInput, Dict[str, torch.Tensor]]:
    """ Get the driver input and broadcast it to other workers.  """
    assert self.is_driver_worker

    worker_input: WorkerInput = self.prepare_worker_input(
        execute_model_req=execute_model_req)
    model_input: ModelRunnerInputBase = (
        self.model_runner.prepare_model_input(
            execute_model_req.seq_group_metadata_list,
            execute_model_req.virtual_engine,
            execute_model_req.finished_requests_ids))

    kwargs = extract_previous_hidden_states(execute_model_req)

    if self.do_metadata_broadcast:
        broadcast_data = worker_input.as_broadcastable_tensor_dict()
        broadcast_data.update(model_input.as_broadcastable_tensor_dict())
        broadcast_data.update(kwargs)
        '''
        ==========================
        Modify by Context Parallel
        ==========================
        @brief: add seq_group metadata to broadcast.
        '''
        broadcast_data['seq_group_metadata'] = execute_model_req.seq_group_metadata_list
        '''
        =======================
        End of Context Parallel
        =======================
        '''
        broadcast_tensor_dict(broadcast_data, src=0)

    if execute_model_req.async_callback:
        model_input = dataclasses.replace(  # type: ignore
            model_input,
            async_callback=execute_model_req.async_callback)

    return model_input, worker_input, kwargs

def vllm__worker__worker_base__LocalOrDistributedWorkerBase___get_worker_input_from_broadcast(
    self
) -> Optional[Tuple[BroadcastableModelInput, WorkerInput, Dict[
        str, torch.Tensor]]]:
    """ Get the worker input from the broadcasted tensor dict. """
    assert self.do_metadata_broadcast
    assert not self.is_driver_worker
    broadcast_data = broadcast_tensor_dict(src=0)
    if not broadcast_data:
        return None

    worker_input = WorkerInput.from_broadcasted_tensor_dict(broadcast_data)
    model_input = (
        self.model_runner.make_model_input_from_broadcasted_tensor_dict(
            broadcast_data))

    kwargs = extract_previous_hidden_states(broadcast_data)

    return model_input, worker_input, kwargs


def vllm__worker__worker_base__LocalOrDistributedWorkerBase__prepare_input(
    self,
    execute_model_req: Optional[ExecuteModelRequest] = None
) -> Optional[Tuple[BroadcastableModelInput, WorkerInput, Dict[str, torch.Tensor]]]:
    """
    Prepare the inputs to ModelRunner and workers.
    """
    if self.is_driver_worker:
        if execute_model_req is None:
            if self.do_metadata_broadcast:
                # This signals that there's no more requests to process for
                # now. All workers are running infinite loop with
                # broadcast_tensor_dict, and it stops the loop when the
                # driver broadcasts an empty input. Send an empty input to
                # notify all other workers to stop their execution loop.
                broadcast_tensor_dict({}, src=0)
            return None
        return self._get_driver_input_and_broadcast(execute_model_req)
    else:
        return self._get_worker_input_from_broadcast()

MluHijackObject.apply_hijack(
    LocalOrDistributedWorkerBase,
    LocalOrDistributedWorkerBase._get_driver_input_and_broadcast,
    vllm__worker__worker_base__LocalOrDistributedWorkerBase___get_driver_input_and_broadcast)

MluHijackObject.apply_hijack(
    LocalOrDistributedWorkerBase,
    LocalOrDistributedWorkerBase._get_worker_input_from_broadcast,
    vllm__worker__worker_base__LocalOrDistributedWorkerBase___get_worker_input_from_broadcast)

MluHijackObject.apply_hijack(
    LocalOrDistributedWorkerBase,
    LocalOrDistributedWorkerBase.prepare_input,
    vllm__worker__worker_base__LocalOrDistributedWorkerBase__prepare_input)
