import torch
from typing import (TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Set,
                    Tuple, Type, TypeVar, Union)
from vllm.forward_context import set_forward_context
from vllm.multimodal.inputs import MultiModalKwargs
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm_mlu._mlu_utils import *
from vllm.worker.model_runner import (
    TModelInputForGPU, ModelInputForGPU,
    ModelInputForGPUWithSamplingMetadata,
    ModelInputForGPUBuilder, GPUModelRunnerBase,
    ModelRunner, CUDAGraphRunner,
    LORA_WARMUP_RANK, _get_graph_batch_size,
    _BATCH_SIZES_TO_CAPTURE, _NUM_WARMUP_ITERS
)
from vllm.worker.mlu_model_runner import MLUModelRunner
from vllm.sequence import (IntermediateTensors, SequenceGroupMetadata)
from vllm.distributed import get_pp_group
from vllm.model_executor.layers.sampler import SamplerOutput
from ..zigzag_utils import get_context_model_parallel_world_size, zigzag_split
import vllm.envs as envs

try:
    from flashinfer import BatchDecodeWithPagedKVCacheWrapper
    from flashinfer.decode import CUDAGraphBatchDecodeWithPagedKVCacheWrapper
    from flashinfer.prefill import BatchPrefillWithPagedKVCacheWrapper
    FLASHINFER_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024
except ImportError:
    BatchDecodeWithPagedKVCacheWrapper = None
    CUDAGraphBatchDecodeWithPagedKVCacheWrapper = None
    BatchPrefillWithPagedKVCacheWrapper = None
    FLASHINFER_WORKSPACE_BUFFER_SIZE = 0

_PAD_SLOT_ID = -1

@torch.inference_mode()
def vllm__worker__mlu_model_runner__MLUModelRunner__execute_model(
    self,
    model_input: ModelInputForGPUWithSamplingMetadata,
    kv_caches: List[torch.Tensor],
    intermediate_tensors: Optional[IntermediateTensors] = None,
    num_steps: int = 1,
) -> Optional[Union[List[SamplerOutput], IntermediateTensors]]:
    if num_steps > 1:
        raise ValueError("num_steps > 1 is not supported in ModelRunner")

    if self.lora_config:
        assert model_input.lora_requests is not None
        assert model_input.lora_mapping is not None
        self.set_active_loras(model_input.lora_requests,
                              model_input.lora_mapping)

    if self.prompt_adapter_config:
        assert model_input.prompt_adapter_requests is not None
        assert model_input.prompt_adapter_mapping is not None
        self.set_active_prompt_adapters(
            model_input.prompt_adapter_requests,
            model_input.prompt_adapter_mapping)

    self.attn_state.begin_forward(model_input)

    # Currently cuda graph is only supported by the decode phase.
    assert model_input.attn_metadata is not None
    prefill_meta = model_input.attn_metadata.prefill_metadata
    decode_meta = model_input.attn_metadata.decode_metadata
    # TODO(andoorve): We can remove this once all
    # virtual engines share the same kv cache.
    virtual_engine = model_input.virtual_engine
    if prefill_meta is None and decode_meta.use_cuda_graph:
        assert model_input.input_tokens is not None
        graph_batch_size = model_input.input_tokens.shape[0]
        model_executable = self.graph_runners[virtual_engine][
            graph_batch_size]
    else:
        model_executable = self.model

    multi_modal_kwargs = model_input.multi_modal_kwargs or {}
    seqlen_agnostic_kwargs = {
        "finished_requests_ids": model_input.finished_requests_ids,
        "request_ids_to_seq_ids": model_input.request_ids_to_seq_ids,
    } if self.has_inner_state else {}
    if (self.observability_config is not None
            and self.observability_config.collect_model_forward_time):
        model_forward_start = torch.mlu.Event(enable_timing=True)
        model_forward_end = torch.mlu.Event(enable_timing=True)
        model_forward_start.record()

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: add mlu metrics
    '''
    # Add time markers for model_executable+compute_logits
    if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
        use_cuda_graph = ((prefill_meta is None and decode_meta.use_cuda_graph)
                          or use_context_mlugraph)
        # if use_cuda_graph, the start timestamp will be inserted inside MLUGraphRunner.forward()
        if not use_cuda_graph:
            start = torch.mlu.Event(enable_timing=True)
            start.record()

    '''
    ==========================
    Modify by Context Parallel
    ==========================
    @brief: context parallel split input for model with zigzag method
    '''
    if get_context_model_parallel_world_size() > 1 and model_input.attn_metadata.prefill_metadata:
        with set_forward_context(model_input.attn_metadata):
            zigzag_input_ids, zigzag_positions, zigzag_attn_metadata = zigzag_split(model_input.input_tokens,
                                                                                    model_input.input_positions,
                                                                                    model_input.attn_metadata, _PAD_SLOT_ID)
            hidden_or_intermediate_states = model_executable(
                input_ids=zigzag_input_ids,
                positions=zigzag_positions,
                kv_caches=kv_caches,
                attn_metadata=zigzag_attn_metadata,
                intermediate_tensors=intermediate_tensors,
                **multi_modal_kwargs,
                **seqlen_agnostic_kwargs)
    else:
        with set_forward_context(model_input.attn_metadata):
            hidden_or_intermediate_states = model_executable(
                input_ids=model_input.input_tokens,
                positions=model_input.input_positions,
                kv_caches=kv_caches,
                attn_metadata=model_input.attn_metadata,
                intermediate_tensors=intermediate_tensors,
                **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
                                             device=self.device),
                **seqlen_agnostic_kwargs)

    #################################################################################################
    #                                           DEBUG                                               #
    #################################################################################################
    # import os
    # from vllm.distributed import get_tensor_model_parallel_world_size, get_tensor_model_parallel_rank
    # from from examples.cambricon_custom_funcvllm.mlu_hijack.distributed.parallel_state import (
    #        get_context_model_parallel_rank)
    # from ..zigzag_utils import context_parallel_tensor_all_gather, diff1
    # if get_context_model_parallel_world_size() > 1 and attn_metadata.prefill_metadata:
    #     hidden_states = context_parallel_tensor_all_gather(hidden_states, zigzag_attn_metadata, dim=0)
    # if attn_metadata.prefill_metadata and (kv_caches[0] is not None):
    #     file_path = '/workspace/output_base_' + str(hidden_states.shape) + \
    #                 '_tp_' + str(get_tensor_model_parallel_world_size()) + '.pth'
    #     if get_context_model_parallel_rank() == 0 and get_tensor_model_parallel_rank() == 0:
    #         if os.path.exists(file_path):
    #             print("##################compare################")
    #             hidden_states_base = torch.load(file_path)
    #             print("########output_diff1: ", diff1(hidden_states, hidden_states_base))
    #         else:
    #             print("##################save base################")
    #             torch.save(hidden_states, file_path)

    '''
    @brief: logits_processor in context parallel need attn_metadata param
    '''
    if get_context_model_parallel_world_size() > 1 and model_input.attn_metadata.prefill_metadata:
        setattr(self.model.logits_processor, 'attn_metadata', zigzag_attn_metadata)
    else:
        setattr(self.model.logits_processor, 'attn_metadata', None)
    '''
    =======================
    End of Context Parallel
    =======================
    '''

    if (self.observability_config is not None
            and self.observability_config.collect_model_forward_time):
        model_forward_end.record()

    # Compute the logits in the last pipeline stage.
    if not get_pp_group().is_last_rank:
        if (self.is_driver_worker
                and hidden_or_intermediate_states is not None
                and isinstance(hidden_or_intermediate_states,
                               IntermediateTensors)
                and self.observability_config is not None
                and self.observability_config.collect_model_forward_time):
            model_forward_end.synchronize()
            model_forward_time = model_forward_start.elapsed_time(
                model_forward_end)
            orig_model_forward_time = 0.0
            if intermediate_tensors is not None:
                orig_model_forward_time = intermediate_tensors.tensors.get(
                    "model_forward_time", torch.tensor(0.0)).item()
            hidden_or_intermediate_states.tensors["model_forward_time"] = (
                torch.tensor(model_forward_time + orig_model_forward_time))
        return hidden_or_intermediate_states

    logits = self.model.compute_logits(hidden_or_intermediate_states,
                                       model_input.sampling_metadata)

    # Add time markers for model_executable+compute_logits
    if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
        end_marker = torch.mlu.Event(enable_timing=True)
        end_marker.record()
        if use_cuda_graph:
            self.time_markers = (model_executable.start, end_marker)
        else:
            self.time_markers = (start, end_marker)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    if not self.is_driver_worker:
        return []

    if model_input.async_callback is not None:
        model_input.async_callback()

    # Sample the next token.
    output: SamplerOutput = self.model.sample(
        logits=logits,
        sampling_metadata=model_input.sampling_metadata,
    )
    if (self.observability_config is not None
            and self.observability_config.collect_model_forward_time
            and output is not None):
        model_forward_end.synchronize()
        model_forward_time = model_forward_start.elapsed_time(
            model_forward_end)
        orig_model_forward_time = 0.0
        if intermediate_tensors is not None:
            orig_model_forward_time = intermediate_tensors.tensors.get(
                "model_forward_time", torch.tensor(0.0)).item()
        # If there are multiple workers, we are still tracking the latency
        # from the start time of the driver worker to the end time of the
        # driver worker. The model forward time will then end up covering
        # the communication time as well.
        output.model_forward_time = (orig_model_forward_time +
                                     model_forward_time)


    if self.return_hidden_states:
        # we only need to pass hidden states of most recent token
        assert model_input.sampling_metadata is not None
        indices = model_input.sampling_metadata.selected_token_indices
        if model_input.is_prompt:
            hidden_states = hidden_or_intermediate_states.index_select(
                0, indices)
        elif decode_meta.use_cuda_graph:
            hidden_states = hidden_or_intermediate_states[:len(indices)]
        else:
            hidden_states = hidden_or_intermediate_states

        output.hidden_states = hidden_states

    return [output]


MluHijackObject.apply_hijack(MLUModelRunner,
                             MLUModelRunner.execute_model,
                             vllm__worker__mlu_model_runner__MLUModelRunner__execute_model)
