import time
from typing import Optional, List, Union, Mapping
from vllm.engine.llm_engine import LLMEngine
from vllm_mlu._mlu_utils import USE_PAGED, BlockSizeInfo
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.sampling_params import SamplingParams
from vllm.lora.request import LoRARequest
from vllm.logger import init_logger
from vllm.inputs import PromptType
from vllm.pooling_params import PoolingParams
from vllm.prompt_adapter.request import PromptAdapterRequest
from vllm.utils import deprecate_kwargs

logger = init_logger(__name__)


@deprecate_kwargs(
    "inputs",
    additional_message="Please use the 'prompt' parameter instead.",
)
def vllm_engine__llm_engine__LLMEngine__add_request(
        self,
        request_id: str,
        prompt: Optional[PromptType] = None,
        params: Optional[Union[SamplingParams, PoolingParams]] = None,
        arrival_time: Optional[float] = None,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
        *,
        inputs: Optional[PromptType] = None,  # DEPRECATED
) -> None:
    """Add a request to the engine's request pool.

    The request is added to the request pool and will be processed by the
    scheduler as `engine.step()` is called. The exact scheduling policy is
    determined by the scheduler.

    Args:
        request_id: The unique ID of the request.
        prompt: The prompt to the LLM. See :class:`~vllm.inputs.PromptType`
            for more details about the format of each input.
        params: Parameters for sampling or pooling.
            :class:`~vllm.SamplingParams` for text generation.
            :class:`~vllm.PoolingParams` for pooling.
        arrival_time: The arrival time of the request. If None, we use
            the current monotonic time.
        trace_headers: OpenTelemetry trace headers.
        priority: The priority of the request.
            Only applicable with priority scheduling.

    Details:
        - Set arrival_time to the current time if it is None.
        - Set prompt_token_ids to the encoded prompt if it is None.
        - Create `n` number of :class:`~vllm.Sequence` objects.
        - Create a :class:`~vllm.SequenceGroup` object
            from the list of :class:`~vllm.Sequence`.
        - Add the :class:`~vllm.SequenceGroup` object to the scheduler.

    Example:
        >>> # initialize engine
        >>> engine = LLMEngine.from_engine_args(engine_args)
        >>> # set request arguments
        >>> example_prompt = "Who is the president of the United States?"
        >>> sampling_params = SamplingParams(temperature=0.0)
        >>> request_id = 0
        >>>
        >>> # add the request to the engine
        >>> engine.add_request(
        >>>    str(request_id),
        >>>    example_prompt,
        >>>    SamplingParams(temperature=0.0))
        >>> # continue the request processing
        >>> ...
    """
    if inputs is not None:
        prompt = inputs
    assert prompt is not None and params is not None

    if lora_request is not None and not self.lora_config:
        raise ValueError(f"Got lora_request {lora_request} but LoRA is "
                            "not enabled!")

    if priority != 0 and not self.scheduler_config.policy == "priority":
        raise ValueError(f"Got priority {priority} but "
                            "Priority scheduling is not enabled.")

    if isinstance(params, SamplingParams) \
        and (params.guided_decoding or params.logits_processors) \
        and self.scheduler_config.num_scheduler_steps > 1:
        raise ValueError(
            "Guided decoding and logits processors are not supported "
            "in multi-step decoding")

    if arrival_time is None:
        arrival_time = time.time()

    if self.tokenizer is not None:
        self._validate_token_prompt(
            prompt,
            tokenizer=self.get_tokenizer(lora_request=lora_request))

    preprocessed_inputs = self.input_preprocessor.preprocess(
        prompt,
        request_id=request_id,
        lora_request=lora_request,
        prompt_adapter_request=prompt_adapter_request,
    )
    processed_inputs = self.input_processor(preprocessed_inputs)

    '''
    =============================
    Added by vllm_mlu
    =============================
    @brief: check input_len + output_len <= block_size
    '''
    def check_block_size_valid(input_len, output_len):
        if BlockSizeInfo.BLOCK_SIZE < input_len + output_len:
            raise ValueError(f"BLOCK_SIZE({BlockSizeInfo.BLOCK_SIZE}) can't smaller than " +
                             f"input_len({input_len}) + output_len({output_len}).")

    if isinstance(params, SamplingParams):
        output_len = params.max_tokens

        # Check for 'prompt_token_ids' in different levels of processed_inputs
        if not USE_PAGED:
            for key in ['prompt_token_ids', 'encoder', 'decoder']:
                if key in processed_inputs:
                    if key == 'prompt_token_ids':
                        input_len = len(processed_inputs[key])
                    elif isinstance(processed_inputs[key], dict) and 'prompt_token_ids' in processed_inputs[key]:
                        input_len = len(processed_inputs[key]['prompt_token_ids'])
                    else:
                        continue

                    check_block_size_valid(input_len, output_len)
    '''
    ==================
    End of modification
    ==================
    '''

    self._add_processed_request(
        request_id=request_id,
        processed_inputs=processed_inputs,
        params=params,
        arrival_time=arrival_time,
        lora_request=lora_request,
        prompt_adapter_request=prompt_adapter_request,
        trace_headers=trace_headers,
        priority=priority,
    )


def vllm__engine__llm_engine__LLMEngine__get_latency(self):
    latency = self.model_executor.get_latency()
    return latency


def vllm__engine__llm_engine__LLMEngine__get_memory_usage(self):
    return self.model_executor.get_memory_usage()


def vllm__engine__llm_engine__LLMEngine__get_block_usage(self):
    assert len(self.scheduler) == 1, f"Only support pipeline_parallel_size=1."
    num_free_gpu_blocks = self.scheduler[0].block_manager.get_num_free_gpu_blocks()
    num_free_cpu_blocks = self.scheduler[0].block_manager.get_num_free_cpu_blocks()
    return (num_free_gpu_blocks, num_free_cpu_blocks)


# for client init/reset server scheduler profile data
def vllm__engine__llm_engine__LLMEngine__init_scheduler_view(self):
    for scheduler in self.scheduler:
        if hasattr(scheduler, "init_scheduler_view"):
            scheduler.init_scheduler_view()
        else:
            logger.warning("Can not find any scheduler view, " +
                           "please 'export VLLM_SCHEDULER_PROFILE=true' first.")


# for client pulling server scheduler profile data
def vllm__engine__llm_engine__LLMEngine__save_scheduler_view(self):
    for idx, scheduler in enumerate(self.scheduler):
        if hasattr(scheduler, "save_scheduler_view"):
            scheduler.save_scheduler_view(idx)
        else:
            logger.warning("Can not find any scheduler view, " +
                           "please 'export VLLM_SCHEDULER_PROFILE=true' first.")


MluHijackObject.apply_hijack(LLMEngine,
                             "init_scheduler_view",
                             vllm__engine__llm_engine__LLMEngine__init_scheduler_view)
MluHijackObject.apply_hijack(LLMEngine,
                             "save_scheduler_view",
                             vllm__engine__llm_engine__LLMEngine__save_scheduler_view)
MluHijackObject.apply_hijack(LLMEngine,
                             LLMEngine.add_request,
                             vllm_engine__llm_engine__LLMEngine__add_request)
MluHijackObject.apply_hijack(LLMEngine,
                             "get_latency",
                             vllm__engine__llm_engine__LLMEngine__get_latency)
MluHijackObject.apply_hijack(LLMEngine,
                             "get_memory_usage",
                             vllm__engine__llm_engine__LLMEngine__get_memory_usage)
MluHijackObject.apply_hijack(LLMEngine,
                             "get_block_usage",
                             vllm__engine__llm_engine__LLMEngine__get_block_usage)
