import asyncio
from abc import abstractmethod
from typing import List, Optional

from vllm.executor.distributed_gpu_executor import DistributedGPUExecutorAsync
from vllm.sequence import ExecuteModelRequest
from vllm.model_executor.layers.sampler import SamplerOutput
from vllm_mlu.mlu_hijack_utils import MluHijackObject

from examples.cambricon_custom_func.vllm.mlu_hijack.common import init_logger

logger = init_logger(__name__)

'''
=============================
Modify by vllm_mlu
=============================
@brief: Add two parameters, in which prefill_locs indicates the start location
and token_chunk_sizes indicates the chunk size for each task.
'''
async def vllm__executor__distributed_gpu_executor__DistributedGPUExecutorAsync__execute_model_async(
        self,
        execute_model_req: ExecuteModelRequest,
        prefill_locs: Optional[List[int]] = None,
        token_chunk_sizes: Optional[List[int]] = None,
) -> List[SamplerOutput]:
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    if self.parallel_worker_tasks is None:
        # Start model execution loop running in the parallel workers
        self.parallel_worker_tasks = asyncio.create_task(
            self._start_worker_execution_loop())

    # Only the driver worker returns the sampling results.
    return await self._driver_execute_model_async(execute_model_req, prefill_locs, token_chunk_sizes)

'''
=============================
Modify by vllm_mlu
=============================
@brief: Add two parameters, in which prefill_locs indicates the start location
and token_chunk_sizes indicates the chunk size for each task.
'''
@abstractmethod
async def vllm__executor__distributed_gpu_executor__DistributedGPUExecutorAsync___driver_execute_model_async(
    self,
    execute_model_req: Optional[ExecuteModelRequest] = None,
    prefill_locs: Optional[List[int]] = None,
    token_chunk_sizes: Optional[List[int]] = None,
) -> List[SamplerOutput]:
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    """Execute the model asynchronously in the driver worker.

    Passing None will cause the driver to stop the model execution
    loop running in each of the remote workers.
    """
    raise NotImplementedError

MluHijackObject.apply_hijack(
    DistributedGPUExecutorAsync,
    DistributedGPUExecutorAsync.execute_model_async,
    vllm__executor__distributed_gpu_executor__DistributedGPUExecutorAsync__execute_model_async
)
MluHijackObject.apply_hijack(
    DistributedGPUExecutorAsync,
    DistributedGPUExecutorAsync._driver_execute_model_async,
    vllm__executor__distributed_gpu_executor__DistributedGPUExecutorAsync___driver_execute_model_async
)
