import dataclasses
import importlib
import os
import time
from abc import ABC, abstractmethod
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union

import torch

from vllm.distributed import broadcast_tensor_dict, get_pp_group, get_tp_group
from vllm.logger import init_logger
from vllm.lora.request import LoRARequest
from vllm.model_executor.layers.sampler import SamplerOutput
from vllm.platforms import current_platform
from vllm.sequence import (ExecuteModelRequest, IntermediateTensors)
from vllm.utils import (enable_trace_function_call_for_thread,
                        update_environment_variables)
from vllm.worker.model_runner_base import (BroadcastableModelInput,
                                           ModelRunnerBase,
                                           ModelRunnerInputBase)
from vllm.worker.worker_base import (LocalOrDistributedWorkerBase,
                                     WorkerInput,
                                     extract_previous_hidden_states)
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from examples.cambricon_custom_func.vllm.mlu_hijack.common import init_logger

logger = init_logger(__name__)


def vllm__worker__worker_base__LocalOrDistributedWorkerBase___get_driver_input_and_broadcast(
    self, execute_model_req: ExecuteModelRequest,
    prefill_locs: Optional[List[int]] = None,
    token_chunk_sizes: Optional[int] = None,
) -> 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)
    """
    ======================================
    Modified by Chunked Parallel Pipeline.
    ======================================
    Pass prefill location and chunk size parameters.
    """
    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,
            prefill_locs,
            token_chunk_sizes))
    """
    ======================================
    End by Chunked Parallel Pipeline.
    ======================================
    """

    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)
        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__prepare_input(
    self,
    execute_model_req: Optional[ExecuteModelRequest] = None,
    prefill_locs: Optional[List[int]] = None,
    token_chunk_sizes: Optional[int] = 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
        """
        ======================================
        Modified by Chunked Parallel Pipeline.
        ======================================
        Pass prefill location and chunk size parameters.
        """
        return self._get_driver_input_and_broadcast(
            execute_model_req, prefill_locs, token_chunk_sizes)
        """
        ======================================
        End by Chunked Parallel Pipeline.
        ======================================
        """
    else:
        return self._get_worker_input_from_broadcast()

def vllm__worker__worker_base__LocalOrDistributedWorkerBase__execute_model(
    self,
    execute_model_req: Optional[ExecuteModelRequest] = None,
    prefill_locs: Optional[List[int]] = None,
    token_chunk_sizes: Optional[int] = None,
    priority: int = -1,
) -> Optional[List[SamplerOutput]]:
    """Executes at least one model step on the given sequences, unless no
    sequences are provided."""
    start_time = time.perf_counter()

    """
    ======================================
    Modified by Chunked Parallel Pipeline.
    ======================================
    Pass prefill location and chunk size parameters.
    """
    inputs = self.prepare_input(execute_model_req, prefill_locs, token_chunk_sizes)
    """
    ======================================
    End by Chunked Parallel Pipeline.
    ======================================
    """
    if inputs is None:
        return None

    model_input, worker_input, kwargs = inputs
    num_steps = worker_input.num_steps

    self.execute_worker(worker_input)

    # If there is no input, we don't need to execute the model.
    if worker_input.num_seq_groups == 0:
        return []

    """
    ======================================
    Modified by Chunked Parallel Pipeline.
    ======================================
    @brief: To prevent the execution of mlu pipeline interrupted by host communication,
    cancel the host communication and prepare metadata list directly.
    """
    assert (token_chunk_sizes is not None and len(token_chunk_sizes) == 1)
    batch_size = token_chunk_sizes[0]
    metadata_list = self.model_runner.model.get_intermediate_tensor_metadata(
        batch_size,
        dtype=self.model_runner.model_config.dtype,
        device=self.model_runner.device)

    intermediate_tensors = None
    orig_model_execute_time = 0.0
    if not get_pp_group().is_first_rank:
        intermediate_tensors = IntermediateTensors(
            get_pp_group().recv_tensor_dict(
                all_gather_group=get_tp_group(),
                recv_metadata_list=metadata_list))
        if (self.observability_config is not None
                and self.observability_config.collect_model_execute_time):
            orig_model_execute_time = intermediate_tensors.tensors.get(
                "model_execute_time", torch.tensor(0)).item()
    """
    ======================================
    End by Chunked Parallel Pipeline.
    ======================================
    """

    output = self.model_runner.execute_model(
        model_input=model_input,
        kv_caches=self.kv_cache[worker_input.virtual_engine]
        if self.kv_cache is not None else None,
        intermediate_tensors=intermediate_tensors,
        num_steps=num_steps,
        **kwargs,
    )

    model_execute_time = time.perf_counter() - start_time
    if not get_pp_group().is_last_rank:
        # output is IntermediateTensors
        if (self.observability_config is not None
                and self.observability_config.collect_model_execute_time):
            output.tensors["model_execute_time"] = torch.tensor(
                model_execute_time + orig_model_execute_time)
        get_pp_group().send_tensor_dict(output.tensors,
                                        all_gather_group=get_tp_group())
        return [None]
    if (self.observability_config is not None
            and self.observability_config.collect_model_execute_time
            and output is not None):
        for o in output:
            o.model_execute_time = (orig_model_execute_time +
                                    model_execute_time)

    # output is List[SamplerOutput]
    return output

MluHijackObject.apply_hijack(
    LocalOrDistributedWorkerBase,
    LocalOrDistributedWorkerBase.prepare_input,
    vllm__worker__worker_base__LocalOrDistributedWorkerBase__prepare_input
)
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.execute_model,
    vllm__worker__worker_base__LocalOrDistributedWorkerBase__execute_model
)
