import weakref
from typing import (List, Optional)

import torch
import torch.distributed

from vllm.compilation.compile_context import set_compile_context
from vllm.distributed import get_pp_group
from vllm.inputs import INPUT_REGISTRY
from vllm.lora.request import LoRARequest
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.sampling_params import SamplingParams
from vllm.sequence import SequenceGroupMetadata
from vllm.worker.model_runner import (
    TModelInputForGPU,
    LORA_WARMUP_RANK,
    _BATCH_SIZES_TO_CAPTURE
)
from vllm.worker.mlu_model_runner import (
    MLUModelRunnerBase,
    ModelInputForMLUBuilder
)
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from examples.cambricon_custom_func.vllm.mlu_hijack.common import init_logger

logger = init_logger(__name__)


@torch.inference_mode()
def vllm__worker__mlu_model_runner__MLUModelRunnerBase__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
    # This represents the maximum number of different requests
    # that will have unique loras, an therefore the max amount of memory
    # consumption create dummy lora request copies from the lora request
    # passed in, which contains a lora from the lora warmup path.
    dummy_lora_requests: List[LoRARequest] = []
    dummy_lora_requests_per_seq: List[LoRARequest] = []
    if self.lora_config:
        assert self.lora_manager is not None
        with self.lora_manager.dummy_lora_cache():
            for idx in range(self.lora_config.max_loras):
                lora_id = idx + 1
                dummy_lora_request = LoRARequest(
                    lora_name=f"warmup_{lora_id}",
                    lora_int_id=lora_id,
                    lora_path="/not/a/real/path",
                )
                self.lora_manager.add_dummy_lora(dummy_lora_request,
                                                    rank=LORA_WARMUP_RANK)
                dummy_lora_requests.append(dummy_lora_request)
            dummy_lora_requests_per_seq = [
                dummy_lora_requests[idx % len(dummy_lora_requests)]
                for idx in range(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] = []
    # Additional GPU memory may be needed for multi-modal encoding, which
    # needs to be accounted for when calculating the GPU blocks for
    # vLLM blocker manager.
    # To exercise the worst scenario for GPU memory consumption,
    # the number of seqs (batch_size) is chosen to maximize the number
    # of images processed.

    max_mm_tokens = self.mm_registry.get_max_multimodal_tokens(
        self.model_config)
    if max_mm_tokens > 0:
        max_num_seqs_orig = max_num_seqs
        max_num_seqs = min(max_num_seqs,
                            max_num_batched_tokens // max_mm_tokens)
        if max_num_seqs < 1:
            expr = (f"min({max_num_seqs_orig}, "
                    f"{max_num_batched_tokens} // {max_mm_tokens})")
            logger.warning(
                "Computed max_num_seqs (%s) to be less than 1. "
                "Setting it to the minimum value of 1.", expr)
            max_num_seqs = 1

    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

        dummy_data = self.input_registry \
            .dummy_data_for_profiling(self.model_config,
                                        seq_len,
                                        self.mm_registry)

        seq = SequenceGroupMetadata(
            request_id=str(group_id),
            is_prompt=True,
            seq_data={group_id: dummy_data.seq_data},
            sampling_params=sampling_params,
            block_tables=None,
            lora_request=dummy_lora_requests_per_seq[group_id]
            if dummy_lora_requests_per_seq else None,
            multi_modal_data=dummy_data.multi_modal_data,
            multi_modal_placeholders=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).
    # it is important to create tensors inside the loop, rather than
    # multiplying the list, to avoid Dynamo from treating them as
    # tensor aliasing.
    '''
    =============================
    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]
    """
    ======================================
    Modified by Chunked Parallel Pipeline.
    ======================================
    @brief: Add two parameters: prefill_loc and token_chunk_size.
    """
    token_chunk_sizes = [seq.token_chunk_size for seq in seqs]
    model_input = self.prepare_model_input(
        seqs,
        finished_requests_ids=finished_requests_ids,
        prefill_locs=[0]*len(seqs),
        token_chunk_sizes=token_chunk_sizes,
    )
    """
    ======================================
    End by Chunked Parallel Pipeline.
    ======================================
    """
    intermediate_tensors = None
    if not get_pp_group().is_first_rank:
        intermediate_tensors = self.model.make_empty_intermediate_tensors(
            batch_size=batch_size,
            dtype=self.model_config.dtype,
            device=self.device)

    graph_batch_size = self.max_batchsize_to_capture
    batch_size_capture_list = [
        bs for bs in _BATCH_SIZES_TO_CAPTURE if bs <= graph_batch_size
    ]
    if self.model_config.enforce_eager:
        batch_size_capture_list = []
    with set_compile_context(batch_size_capture_list):
        self.execute_model(model_input, kv_caches, intermediate_tensors)
    torch.mlu.synchronize()

    return


MluHijackObject.apply_hijack(
    MLUModelRunnerBase,
    MLUModelRunnerBase.profile_run,
    vllm__worker__mlu_model_runner__MLUModelRunnerBase__profile_run
)
