from typing import Optional, Dict, Any
from vllm.entrypoints.llm import LLM
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.logger import init_logger
from vllm.engine.arg_utils import (EngineArgs, HfOverrides, PoolerConfig,
                                   TaskOption)


logger = init_logger(__name__)


vllm__entrypoints__llm__LLM____init__org = LLM.__init__

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: add cp and ep parameter
    '''
    # pop context_parallel_size
    EngineArgs.context_parallel_size = kwargs.pop("context_parallel_size", 1)
    # pop moe_tp_size and moe_ep_size
    EngineArgs.moe_tp_size = kwargs.pop("moe_tp_size", -1)
    # pop moe_ep_size and moe_ep_size
    EngineArgs.moe_ep_size = kwargs.pop("moe_ep_size", -1)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    vllm__entrypoints__llm__LLM____init__org(
        self=self,
        model=model,
        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,
        # After positional args are removed, move this right below `model`
        task=task,
        override_pooler_config=override_pooler_config,
        **kwargs
    )


MluHijackObject.apply_hijack(LLM,
                             LLM.__init__,
                             vllm__entrypoints__llm__LLM____init__)
