import dataclasses
import weakref
from typing import (List, Optional, TypeVar)

from vllm.distributed import get_pp_group
from vllm.model_executor import SamplingMetadata
from vllm.sequence import SequenceGroupMetadata
from vllm.worker.model_runner import (
    GPUModelRunnerBase,
    ModelInputForGPUBuilder,
    ModelInputForGPUWithSamplingMetadata,
    ModelRunner,
    TModelInputForGPU
)
from vllm_mlu.mlu_hijack_utils import MluHijackObject

from examples.cambricon_custom_func.vllm.mlu_hijack.common import init_logger

logger = init_logger(__name__)


"""
======================================
Modified by Chunked Parallel Pipeline.
======================================
@brief: Add two parameters, prefill_loc and token_chunk_size.
"""
def vllm__worker__model_runner__ModelInputForGPUBuilder___compute_lens(
    self, inter_data: ModelInputForGPUBuilder.InterDataForSeqGroup,
    seq_idx: int, seq_group_metadata: SequenceGroupMetadata,
    prefill_loc: Optional[int] = None,
    token_chunk_size: Optional[int] = None,
):
    """
    ======================================
    End by Chunked Parallel Pipeline.
    ======================================
    """
    """Compute context length, sequence length and tokens
    for the given sequence data.
    """
    seq_data = seq_group_metadata.seq_data[inter_data.seq_ids[seq_idx]]
    if token_chunk_size is None:
        token_chunk_size = seq_group_metadata.token_chunk_size

    # Compute context length (the number of tokens that are
    # already computed) and sequence length (total number of tokens).

    seq_len = seq_data.get_len()
    if inter_data.is_prompt:
        """
        ======================================
        Modified by Chunked Parallel Pipeline.
        ======================================
        @brief: For chunked pipeline parallel, since multiple tasks
        use the same sequence data with different prefill location,
        an extra parameter is provided to indicate the prefill location.
        """
        context_len = (
            prefill_loc if prefill_loc is not None
            else seq_data.get_num_computed_tokens()
        )
        """
        ======================================
        End by Chunked Parallel Pipeline.
        ======================================
        """
        seq_len = min(seq_len, context_len + token_chunk_size)
    elif self.runner.scheduler_config.is_multi_step or \
        self.runner.model_config.is_encoder_decoder:
        assert prefill_loc is None, "Chunked Parallel Pipeline does not support multi-step."
        context_len = seq_len - 1
    else:
        context_len = seq_data.get_num_computed_tokens()

    # Compute tokens.
    tokens = seq_data.get_token_ids()[context_len:seq_len]

    inter_data.seq_lens[seq_idx] = seq_len
    inter_data.orig_seq_lens[seq_idx] = seq_len
    inter_data.context_lens[seq_idx] = context_len
    inter_data.input_tokens[seq_idx].extend(tokens)
    inter_data.input_positions[seq_idx].extend(range(context_len, seq_len))
    inter_data.query_lens[seq_idx] = seq_len - context_len

    if seq_data.mrope_position_delta is not None:
        if inter_data.mrope_input_positions is None:
            inter_data.mrope_input_positions = [None] * inter_data.n_seqs

        inter_data.mrope_input_positions[
            seq_idx] = MRotaryEmbedding.get_next_input_positions(
                seq_data.mrope_position_delta,
                context_len,
                seq_len,
            )


"""
======================================
Modified by Chunked Parallel Pipeline.
======================================
@brief: Add two parameters, prefill_loc and token_chunk_size.
"""
def vllm__worker__model_runner__ModelInputForGPUBuilder__add_seq_group(
    self, seq_group_metadata: SequenceGroupMetadata,
    prefill_loc: Optional[int] = None,
    token_chunk_size: Optional[int] = None,
):
    """
    ======================================
    End by Chunked Parallel Pipeline.
    ======================================
    """
    """Add a sequence group to the builder."""
    seq_ids = seq_group_metadata.seq_data.keys()
    n_seqs = len(seq_ids)
    is_prompt = seq_group_metadata.is_prompt

    if is_prompt:
        assert n_seqs == 1
        self.decode_only = False

    encoder_seq_len = 0

    if self.runner.model_config.is_encoder_decoder:
        encoder_seq_len = seq_group_metadata.encoder_seq_data.get_len()


    inter_data = self.init_cached_inter_data(
        request_id=seq_group_metadata.request_id,
        seq_ids=seq_ids,
        is_prompt=is_prompt,
        block_tables=seq_group_metadata.block_tables,
        computed_block_nums=seq_group_metadata.computed_block_nums,
        reinit=True,
        reinit_use_defaults=True,
        encoder_seq_len=encoder_seq_len)

    self.inter_data_list.append(inter_data)

    for seq_idx in range(n_seqs):
        for per_seq_fn in self.per_seq_compute_fns:
            """
            ======================================
            Modified by Chunked Parallel Pipeline.
            ======================================
            @brief: Add prefill location and token chunk size parameters.
            """
            if per_seq_fn.__qualname__ == \
                "vllm__worker__model_runner__ModelInputForGPUBuilder___compute_lens":
                per_seq_fn(inter_data, seq_idx, seq_group_metadata, prefill_loc, token_chunk_size)
            else:
                per_seq_fn(inter_data, seq_idx, seq_group_metadata)
            """
            ======================================
            End by Chunked Parallel Pipeline.
            ======================================
            """
    for per_seq_group_fn in self.per_seq_group_compute_fns:
        per_seq_group_fn(inter_data, seq_group_metadata)


def vllm__worker__model_runner__GPUModelRunnerBase___prepare_model_input_tensors(
    self,
    seq_group_metadata_list: List[SequenceGroupMetadata],
    finished_requests_ids: Optional[List[str]] = None,
    prefill_locs: Optional[List[int]] = None,
    token_chunk_sizes: Optional[List[int]] = None,
) -> TModelInputForGPU:
    """Helper method to prepare the model input based on a given sequence
    group. Prepares metadata needed for the base model forward pass but not
    metadata for possible additional steps, e.g., sampling.

    The API assumes seq_group_metedata_list is sorted by prefill -> decode.

    The result tensors and data structure also batches input in prefill
    -> decode order. For example,

    - input_tokens[:num_prefill_tokens] contains prefill tokens.
    - input_tokens[num_prefill_tokens:] contains decode tokens.

    If cuda graph is required, this API automatically pads inputs.
    """
    builder = self._builder_cls(weakref.proxy(self), finished_requests_ids)
    """
    ======================================
    Modified by Chunked Parallel Pipeline.
    ======================================
    @brief: Add two parameters: prefill_loc and token_chunk_size, and
    check whether they are same as sequence group length or empty.
    """
    if prefill_locs is None:
        prefill_locs = [None] * len(seq_group_metadata_list)

    assert len(prefill_locs) == len(seq_group_metadata_list), \
        "the lengths of prefill locs and seq_group_metadata are different."

    if token_chunk_sizes is None:
        token_chunk_sizes = [None] * len(seq_group_metadata_list)

    assert len(token_chunk_sizes) == len(seq_group_metadata_list), \
        "the lengths of token_chunk_sizes and seq_group_metadata are different."

    for seq_group_metadata, prefill_loc, token_chunk_size in zip(
        seq_group_metadata_list, prefill_locs, token_chunk_sizes
    ):
        builder.add_seq_group(seq_group_metadata, prefill_loc, token_chunk_size)
    """
    ======================================
    End by Chunked Parallel Pipeline.
    ======================================
    """

    builder.reset_cached_inter_data()

    return builder.build()  # type: ignore

"""
======================================
Modified by Chunked Parallel Pipeline.
======================================
@brief: Add two parameters, prefill_loc and token_chunk_size.
"""
def vllm__worker__model_runner__ModelRunner__prepare_model_input(
    self,
    seq_group_metadata_list: List[SequenceGroupMetadata],
    virtual_engine: int = 0,
    finished_requests_ids: Optional[List[str]] = None,
    prefill_locs: Optional[List[int]] = None,
    token_chunk_sizes: Optional[List[int]] = None,
) -> ModelInputForGPUWithSamplingMetadata:
    """
    ======================================
    End by Chunked Parallel Pipeline.
    ======================================
    """
    """Prepare the model input based on a given sequence group, including
    metadata for the sampling step.

    The API assumes seq_group_metadata_list is sorted by prefill -> decode.

    The result tensors and data structure also batches input in prefill
    -> decode order. For example,

    - input_tokens[:num_prefill_tokens] contains prefill tokens.
    - input_tokens[num_prefill_tokens:] contains decode tokens.

    If cuda graph is required, this API automatically pads inputs.
    """
    """
    ======================================
    Modified by Chunked Parallel Pipeline.
    ======================================
    Add prefill location parameter.
    """
    model_input = self._prepare_model_input_tensors(
        seq_group_metadata_list,
        finished_requests_ids,
        prefill_locs,
        token_chunk_sizes)
    """
    ======================================
    End by Chunked Parallel Pipeline.
    ======================================
    """
    if get_pp_group().is_last_rank:
        # Sampling metadata is only required for the final pp group
        generators = self.get_generators(finished_requests_ids)
        sampling_metadata = SamplingMetadata.prepare(
            seq_group_metadata_list, model_input.seq_lens,
            model_input.query_lens, self.device, self.pin_memory,
            generators, self.sampling_metadata_cache)
    else:
        sampling_metadata = None
    is_prompt = (seq_group_metadata_list[0].is_prompt
                 if seq_group_metadata_list else None)
    return dataclasses.replace(model_input,
                               sampling_metadata=sampling_metadata,
                               is_prompt=is_prompt,
                               virtual_engine=virtual_engine)

MluHijackObject.apply_hijack(
    ModelInputForGPUBuilder,
    ModelInputForGPUBuilder._compute_lens,
    vllm__worker__model_runner__ModelInputForGPUBuilder___compute_lens
)

MluHijackObject.apply_hijack(
    ModelInputForGPUBuilder,
    ModelInputForGPUBuilder.add_seq_group,
    vllm__worker__model_runner__ModelInputForGPUBuilder__add_seq_group
)

MluHijackObject.apply_hijack(
    GPUModelRunnerBase,
    GPUModelRunnerBase._prepare_model_input_tensors,
    vllm__worker__model_runner__GPUModelRunnerBase___prepare_model_input_tensors
)

MluHijackObject.apply_hijack(
    ModelRunner,
    ModelRunner.prepare_model_input,
    vllm__worker__model_runner__ModelRunner__prepare_model_input
)