import time
from tqdm import tqdm
from typing import Optional, List, Union, Dict, Any
from vllm.entrypoints.llm import LLM
from vllm.engine.arg_utils import (EngineArgs, HfOverrides, PoolerConfig,
                                   TaskOption)
from vllm.engine.llm_engine import LLMEngine
from vllm.outputs import EmbeddingRequestOutput, RequestOutput
from vllm.usage.usage_lib import UsageContext
from vllm.utils import Counter, deprecate_args
from vllm_mlu._mlu_utils import VLLM_LATENCY_DEBUG_EN, VLLM_LATENCY_DEBUG_WITH_DEVICE_EN
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm_mlu.mlu_metric import LLMMetric
from vllm_mlu.dump_info import LLMDumpInfo
from vllm.logger import init_logger


logger = init_logger(__name__)


@deprecate_args(
    start_index=2,  # Ignore self and model
    is_deprecated=lambda: LLM.DEPRECATE_INIT_POSARGS,
    additional_message=(
        "All positional arguments other than `model` will be "
        "replaced with keyword arguments in an upcoming version."),
)
def vllm__entrypoints__llm__LLM____init__(
    self,
    model: str,
    tokenizer: Optional[str] = None,
    tokenizer_mode: str = "auto",
    skip_tokenizer_init: bool = False,
    trust_remote_code: bool = False,
    allowed_local_media_path: str = "",
    tensor_parallel_size: int = 1,
    dtype: str = "auto",
    quantization: Optional[str] = None,
    revision: Optional[str] = None,
    tokenizer_revision: Optional[str] = None,
    seed: int = 0,
    gpu_memory_utilization: float = 0.9,
    swap_space: float = 4,
    cpu_offload_gb: float = 0,
    enforce_eager: Optional[bool] = None,
    max_seq_len_to_capture: int = 8192,
    disable_custom_all_reduce: bool = False,
    disable_async_output_proc: bool = False,
    hf_overrides: Optional[HfOverrides] = None,
    mm_processor_kwargs: Optional[Dict[str, Any]] = None,
    # After positional args are removed, move this right below `model`
    task: TaskOption = "auto",
    override_pooler_config: Optional[PoolerConfig] = None,
    **kwargs,
) -> None:
    '''
    LLM constructor.

    Note: if enforce_eager is unset (enforce_eager is None)
    it defaults to False.
    '''

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: 1) Initialize LLMDumpInfo
            2) Initialize context mlugraph params
    '''
    LLM.dump_info.init_param(
        tensor_parallel_size=tensor_parallel_size, dtype=dtype,
        kv_cache_dtype=kwargs.get('kv_cache_dtype', 'default_value'),
        quantization=quantization,
        model=model, trust_remote_code=kwargs.get('trust_remote_code', 'default_value'))

    enable_context_mlugraph = kwargs.pop("enable_context_mlugraph", False)
    context_batch_size_to_capture = kwargs.pop("context_batch_size_to_capture", None)
    context_seq_len_to_capture = kwargs.pop("context_seq_len_to_capture", None)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

    if "disable_log_stats" not in kwargs:
        kwargs["disable_log_stats"] = True

    engine_args = EngineArgs(
        model=model,
        task=task,
        tokenizer=tokenizer,
        tokenizer_mode=tokenizer_mode,
        skip_tokenizer_init=skip_tokenizer_init,
        trust_remote_code=trust_remote_code,
        allowed_local_media_path=allowed_local_media_path,
        tensor_parallel_size=tensor_parallel_size,
        dtype=dtype,
        quantization=quantization,
        revision=revision,
        tokenizer_revision=tokenizer_revision,
        seed=seed,
        gpu_memory_utilization=gpu_memory_utilization,
        swap_space=swap_space,
        cpu_offload_gb=cpu_offload_gb,
        enforce_eager=enforce_eager,
        max_seq_len_to_capture=max_seq_len_to_capture,
        disable_custom_all_reduce=disable_custom_all_reduce,
        disable_async_output_proc=disable_async_output_proc,
        hf_overrides=hf_overrides,
        mm_processor_kwargs=mm_processor_kwargs,
        override_pooler_config=override_pooler_config,
        **kwargs,
    )

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: set context mlugraph params for EngineArgs
    '''
    setattr(engine_args, "enable_context_mlugraph", enable_context_mlugraph)
    setattr(engine_args, "context_batch_size_to_capture", context_batch_size_to_capture)
    setattr(engine_args, "context_seq_len_to_capture", context_seq_len_to_capture)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

    # Logic to switch between engines is done at runtime instead of import
    # to avoid import order issues
    self.engine_class = self.get_engine_class()

    # TODO(rob): enable mp by default (issue with fork vs spawn)
    self.llm_engine = self.engine_class.from_engine_args(
        engine_args, usage_context=UsageContext.LLM_CLASS)

    self.request_counter = Counter()

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: Get Cpuinfo member for vllm
    '''
    LLM.dump_info.memory_usage()
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


def vllm__entrypoints__llm__LLM__get_metrics(
    self,
    metrics_idx_start,
    only_average,
    input_len,
    output_len,
    tp_nums,
    quantization,
    dump_info=None,
    show_per_iter=False,
) -> None:
    '''
    @brief:该函数用来打印vLLM调用generate接口过程中代码统计的各项性能指标数据
    @params:
        metrics_idx_start: 考虑存在调用generate接口为warmup过程的情况，
        因此设置该参数可忽略统计[0,metrics_idx_start)之间的数据,默认为0,即所有性能数据有效。
        only_average: True 只打印N次调用generate接口的平均性能 False 打印每次调用generate接口的性能及其均值 若N次性能数据波动较大，需自行排查测试环境是否稳定。
        其余参数:均为模型配置参数
    '''
    if VLLM_LATENCY_DEBUG_EN:
        self.metric.calc_metric(self.llm_engine.model_config.model,
                                self.llm_engine.model_config.dtype,
                                metrics_idx_start, only_average,
                                input_len, output_len, tp_nums,
                                quantization, dump_info, show_per_iter)
    else:
        print("Warnning:please set VLLM_LATENCY_DEBUG=true!")


def vllm__entrypoints__llm__LLM___run_engine(
        self, *, use_tqdm: bool
) -> List[Union[RequestOutput, EmbeddingRequestOutput]]:
    # Initialize tqdm.
    if use_tqdm:
        num_requests = self.llm_engine.get_num_unfinished_requests()
        pbar = tqdm(
            total=num_requests,
            desc="Processed prompts",
            dynamic_ncols=True,
            postfix=(f"est. speed input: {0:.2f} toks/s, "
                        f"output: {0:.2f} toks/s"),
        )

    '''
    =============================
    Added by vllm_mlu
    =============================
    '''
    is_latency_debug = VLLM_LATENCY_DEBUG_EN
    # Record start
    if is_latency_debug:
        total_request_num = self.llm_engine.get_num_unfinished_requests()
        self.dump_info.capture_cpu_info()
        peak_memory, block_memory, num_total_gpu_blocks, num_total_cpu_blocks = \
            self.llm_engine.get_memory_usage()
        self.metric.update_memory_usage(peak_memory, block_memory, num_total_gpu_blocks, num_total_cpu_blocks)
        e2e_start_time = self.metric.get_mlu_cost_time()
    '''
    ==================
    End of addition
    ==================
    '''

    # Run the engine.
    outputs: List[Union[RequestOutput, EmbeddingRequestOutput]] = []
    total_in_toks = 0
    total_out_toks = 0
    while self.llm_engine.has_unfinished_requests():
        '''
        =============================
        Added by vllm_mlu
        =============================
        '''
        if is_latency_debug :
            self.dump_info.memory_usage()
            start_time = self.metric.get_mlu_cost_time()
        '''
        ==================
        End of addition
        ==================
        '''
        step_outputs = self.llm_engine.step()
        '''
        =============================
        Added by vllm_mlu
        =============================
        '''
        if is_latency_debug:
            end_time = self.metric.get_mlu_cost_time()
            step_latency = end_time - start_time
            if len(step_outputs) > 0:
                batch_size = len(step_outputs)
                assert batch_size == total_request_num, \
                    f"LLM has received {total_request_num} requests, but only processed {batch_size} requests in the current step.\n" + \
                    f"If you are running benchmark_latency test, please check if the input is correct.\n" + \
                    f"Otherwise, please set env VLLM_LATENCY_DEBUG=false, then run test again.\n"
            num_free_gpu_blocks, num_free_cpu_blocks = self.llm_engine.get_block_usage()
            self.metric.update_step_block_usage(num_free_gpu_blocks, num_free_cpu_blocks)
            self.metric.update_step_latency(step_latency)
            if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
                self.metric.update_step_latency_device(self.llm_engine.get_latency())
            self.dump_info.memory_usage()
        '''
        ==================
        End of addition
        ==================
        '''
        for output in step_outputs:
            if output.finished:
                outputs.append(output)
                if use_tqdm:
                    if isinstance(output, RequestOutput):
                        # Calculate tokens only for RequestOutput
                        assert output.prompt_token_ids is not None
                        total_in_toks += len(output.prompt_token_ids)
                        in_spd = total_in_toks / pbar.format_dict["elapsed"]
                        total_out_toks += sum(
                            len(stp.token_ids) for stp in output.outputs)
                        out_spd = (total_out_toks /
                                    pbar.format_dict["elapsed"])
                        pbar.postfix = (
                            f"est. speed input: {in_spd:.2f} toks/s, "
                            f"output: {out_spd:.2f} toks/s")
                    pbar.update(1)
    '''
    =============================
    Added by vllm_mlu
    =============================
    '''
    if is_latency_debug:
        e2e_end_time = self.metric.get_mlu_cost_time()
        e2e_latency = e2e_end_time - e2e_start_time
        self.metric.add_metrics(batch_size, e2e_latency)
    '''
    ==================
    End of addition
    ==================
    '''

    if use_tqdm:
        pbar.close()
    # Sort the outputs by request ID.
    # This is necessary because some requests may be finished earlier than
    # its previous requests.
    return sorted(outputs, key=lambda x: int(x.request_id))


LLM.metric = LLMMetric()

LLM.dump_info = LLMDumpInfo()

MluHijackObject.apply_hijack(LLM,
                             LLM.__init__,
                             vllm__entrypoints__llm__LLM____init__)
MluHijackObject.apply_hijack(LLM,
                             "get_metrics",
                             vllm__entrypoints__llm__LLM__get_metrics)
MluHijackObject.apply_hijack(LLM,
                             LLM._run_engine,
                             vllm__entrypoints__llm__LLM___run_engine)
