import itertools
from typing import List, Optional, Tuple

import torch
import torch.distributed

from vllm.attention.backends.abstract import AttentionMetadata
from vllm.attention.backends.utils import PAD_SLOT_ID
from vllm.logger import init_logger
from vllm.sampling_params import SamplingParams
from vllm.sequence import SequenceGroupMetadata
from vllm.worker.enc_dec_model_runner import EncoderDecoderModelInput
from vllm.worker.mlu_enc_dec_model_runner import MLUEncoderDecoderModelRunner
from vllm.worker.model_runner import _get_graph_batch_size
from vllm.utils import make_tensor_with_pad

from vllm_mlu.mlu_hijack_utils import MluHijackObject

logger = init_logger(__name__)


class MLUEncoderDecoderModelRunner_V2(MLUEncoderDecoderModelRunner):

    @torch.inference_mode()
    def profile_run(self) -> None:
        # Enable top-k sampling to reflect the accurate memory usage.
        sampling_params = SamplingParams(top_p=0.99, top_k=self.vocab_size - 1)
        max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
        max_num_seqs = self.scheduler_config.max_num_seqs

        # Profile memory usage with max_num_sequences sequences and the total
        # number of tokens equal to max_num_batched_tokens.
        seqs: List[SequenceGroupMetadata] = []

        max_mm_tokens = self.mm_registry.get_max_multimodal_tokens(
            self.model_config)
        if max_mm_tokens > 0:
            logger.info("Starting profile run for multi-modal models.")

        batch_size = 0
        for group_id in range(max_num_seqs):
            seq_len = (max_num_batched_tokens // max_num_seqs +
                       (group_id < max_num_batched_tokens % max_num_seqs))
            batch_size += seq_len

            decoder_dummy_data = self.input_registry \
                .dummy_data_for_profiling(self.model_config,
                                          seq_len,
                                          self.mm_registry,
                                          is_encoder_data=False)
            encoder_dummy_data \
                = self.input_registry.dummy_data_for_profiling(
                    self.model_config,
                                         seq_len,
                                         self.mm_registry,
                                         is_encoder_data=True)

            # Having more tokens is over-conservative but otherwise fine
            assert len(
                decoder_dummy_data.seq_data.prompt_token_ids
            ) >= seq_len, (
                f"Expected at least {seq_len} dummy tokens for profiling, "
                f"but got: {len(decoder_dummy_data.seq_data.prompt_token_ids)}"
            )

            assert decoder_dummy_data.multi_modal_data is None or \
            encoder_dummy_data.multi_modal_data is None, (
                "Multi-modal data can't be provided in both encoder and decoder"
            )

            seq = SequenceGroupMetadata(
                request_id=str(group_id),
                is_prompt=True,
                seq_data={group_id: decoder_dummy_data.seq_data},
                sampling_params=sampling_params,
                block_tables=None,
                encoder_seq_data=encoder_dummy_data.seq_data,
                cross_block_table=None,
                multi_modal_data=decoder_dummy_data.multi_modal_data
                or encoder_dummy_data.multi_modal_data,
                multi_modal_placeholders=decoder_dummy_data.
                multi_modal_placeholders
                or encoder_dummy_data.multi_modal_placeholders)
            seqs.append(seq)

        # Run the model with the dummy inputs.
        num_layers = self.model_config.get_num_layers(self.parallel_config)
        # use an empty tensor instead of `None`` to force Dynamo to pass
        # it by reference, rather by specializing on the value ``None``.
        # the `dtype` argument does not matter, and we use `float32` as
        # a placeholder (it has wide hardware support).
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: support kv cache int8
        '''
        kv_caches = []
        for _ in range(num_layers):
            kv_cache_ = torch.tensor([], dtype=torch.float32, device=self.device)
            kv_cache_scale_ = torch.tensor([], dtype=torch.float32, device=self.device)
            kv_caches.append([kv_cache_, kv_cache_scale_])
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
        finished_requests_ids = [seq.request_id for seq in seqs]
        model_input = self.prepare_model_input(
            seqs, finished_requests_ids=finished_requests_ids)
        intermediate_tensors = None
        self.execute_model(model_input, kv_caches, intermediate_tensors)
        torch.cuda.synchronize()
        return

    def _prepare_encoder_model_input_tensors(
        self,
        seq_group_metadata_list: List[SequenceGroupMetadata],
        model_input: EncoderDecoderModelInput,
    ) -> Tuple[AttentionMetadata, Optional[torch.Tensor],
               Optional[torch.Tensor]]:
        """Helper method to prepare the encoder- and cross-attn-related
        model inputs based on a given sequence group. These additional inputs
        are used to augment an already-computed `EncoderDecoderModelInput`
        data structure which already has decoder-related model inputs
        populated.

        Sets the following attn_metadata fields:
        * `num_encoder_tokens`
        * `encoder_seq_lens`
        * `encoder_seq_lens_tensor`
        * `max_encoder_seq_len`
        * `cross_slot_mapping`
        * `cross_block_tables`

        Constructs a new model inputs data structure, based on
        (1) the existing fields in the `model_inputs` argument,
        and (2) the following additional fields which are
        computed (or in the case of `attn_metadata`, updated)
        by this function:
        * attn_metadata
        * encoder_input_tokens
        * encoder_input_positions

        Arguments:

        * seq_group_metadata_list: list of sequence groups for which to
                                   compute inputs
        * model_inputs: model inputs data structure with decoder-oriented
                        fields already computed.

        Return:

        * Updated model inputs data structure
        """

        if len(seq_group_metadata_list) == 0:
            return (model_input.attn_metadata, None, None)

        # Since we are not supporting chunked prefill either the entire
        # batch is prefill or it is decode
        is_prompt = seq_group_metadata_list[0].is_prompt

        # Build encoder inputs
        encoder_seq_lens: List[int] = []
        if is_prompt:
            # Prefill phase.
            cross_block_tables = self._empty_int32_tensor().view(
                len(seq_group_metadata_list), -1)

            # Extract input tokens/positions, cross-attention slot-mapping,
            # & seq len from each sequence group metadata
            (
                encoder_input_tokens,
                encoder_input_positions,
                cross_slot_mapping,
            ) = (
                [],
                [],
                [],
            )
            for seq_group_metadata in seq_group_metadata_list:
                # Build seq lens
                seq_len = seq_group_metadata.encoder_seq_data.get_len()
                token_ids = seq_group_metadata.encoder_seq_data.get_token_ids()
                encoder_seq_lens.append(seq_len)

                # Build slot mapping
                is_profile_run = (seq_group_metadata.block_tables is None)
                if is_profile_run:
                    # During memory profiling, the block tables are not
                    # initialized yet. In this case, we just use a dummy
                    # slot mapping.
                    # In embeddings, the block tables are {seq_id: None}.
                    cross_slot_mapping.extend([PAD_SLOT_ID] * seq_len)
                else:
                    for i in range(0, seq_len):
                        block_number = seq_group_metadata.cross_block_table[
                            i // self.block_size]
                        block_offset = i % self.block_size
                        slot = block_number * self.block_size + block_offset
                        cross_slot_mapping.append(slot)

                # Build encoder input tokens
                encoder_input_tokens.extend(token_ids)
                encoder_input_positions.extend(list(range(0, seq_len)))

            # Convert tokens/positions & cross-attention
            # slot-mapping to encoder input tensors
            encoder_input_tokens_tensor = self._list_to_long_tensor(
                encoder_input_tokens)
            encoder_input_positions_tensor = self._list_to_long_tensor(
                encoder_input_positions)
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: since `slot_mapping` parameter in tmo.reshape_paged_cache only
            support int32, change dtype to int32.
            '''
            cross_slot_mapping_tensor = self._list_to_int32_tensor(
                cross_slot_mapping)
            '''
            ==================
            End of MLU Hijack
            ==================
            '''

        else:
            # Decode phase.
            encoder_input_tokens_tensor = self._empty_long_tensor()
            encoder_input_positions_tensor = self._empty_long_tensor()
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: since `slot_mapping` parameter in tmo.reshape_paged_cache only
            support int32, change dtype to int32.
            '''
            cross_slot_mapping_tensor = self._empty_int32_tensor()
            '''
            ==================
            End of MLU Hijack
            ==================
            '''
            # Extract cross-attention block tables &
            # seq len from each sequence group metadata.
            # Cross-attention block tables are empty
            # during vLLM memory profiling.
            cross_block_tables = []
            for seq_group_metadata in seq_group_metadata_list:
                for _ in range(len(seq_group_metadata.seq_data)):
                    encoder_seq_lens.append(
                        seq_group_metadata.encoder_seq_data.get_len())
                    cross_block_table = seq_group_metadata.cross_block_table
                    cross_block_tables.append([] if (
                        cross_block_table is None) else cross_block_table)

            if (model_input.attn_metadata is not None
                    and model_input.attn_metadata.use_cuda_graph):
                # We will be using CUDA graph replay for this decode.
                max_len_of_block_table = self.get_max_block_per_batch()
                batch_size = len(encoder_seq_lens)
                graph_batch_size = _get_graph_batch_size(batch_size)
                assert graph_batch_size >= batch_size
                cuda_graph_pad_size = graph_batch_size - batch_size
                # extend the cross_block_tables and encoder_seq_lens to match
                # the graph_batch_size.
                cross_block_tables.extend([[]
                                           for _ in range(cuda_graph_pad_size)
                                           ])
                encoder_seq_lens.extend(
                    itertools.repeat(1, cuda_graph_pad_size))

            else:
                max_len_of_block_table = max(
                    len(block_table) for block_table in cross_block_tables)

            cross_block_tables = make_tensor_with_pad(
                cross_block_tables,
                max_len=max_len_of_block_table,
                pad=0,
                dtype=torch.int32,
                device=self.device,
            )

        # Compute encoder sequence lengths & encoder
        # sequence starting offset tensors
        max_encoder_seq_len = max(encoder_seq_lens, default=0)
        encoder_seq_lens_tensor = self._list_to_int32_tensor(encoder_seq_lens)
        encoder_seq_start_loc = torch.zeros(encoder_seq_lens_tensor.shape[0] +
                                            1,
                                            dtype=torch.int32,
                                            device=self.device)
        torch.cumsum(encoder_seq_lens_tensor,
                     dim=0,
                     dtype=encoder_seq_start_loc.dtype,
                     out=encoder_seq_start_loc[1:])

        # Update attention metadata with encoder-oriented attributes
        attn_metadata = model_input.attn_metadata
        assert attn_metadata is not None
        (
            attn_metadata.num_encoder_tokens,
            attn_metadata.encoder_seq_lens,
            attn_metadata.encoder_seq_lens_tensor,
            attn_metadata.max_encoder_seq_len,
            attn_metadata.encoder_seq_start_loc,
            attn_metadata.cross_slot_mapping,
            attn_metadata.cross_block_tables,
        ) = (
            sum(encoder_seq_lens),
            encoder_seq_lens,
            encoder_seq_lens_tensor,
            max_encoder_seq_len,
            encoder_seq_start_loc,
            cross_slot_mapping_tensor,
            cross_block_tables,
        )

        return (attn_metadata, encoder_input_tokens_tensor,
                encoder_input_positions_tensor)

MluHijackObject.apply_hijack(MLUEncoderDecoderModelRunner,
                             MLUEncoderDecoderModelRunner.profile_run,
                             MLUEncoderDecoderModelRunner_V2.profile_run)
MluHijackObject.apply_hijack(MLUEncoderDecoderModelRunner,
                             MLUEncoderDecoderModelRunner._prepare_encoder_model_input_tensors,
                             MLUEncoderDecoderModelRunner_V2._prepare_encoder_model_input_tensors)
