import asyncio
from typing import List, Optional

from vllm.executor.distributed_mlu_executor import DistributedMLUExecutorAsync
from vllm.executor.ray_mlu_executor import RayMLUExecutorAsync
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.model_executor.layers.sampler import SamplerOutput
from vllm.sequence import ExecuteModelRequest

from examples.cambricon_custom_func.vllm.mlu_hijack.common import init_logger
from ..lock_utils import (_run_task_with_priority_lock, PriorityLock)

logger = init_logger(__name__)

vllm__executor__ray_mlu_executor__RayMLUExecutorAsync____init____org = RayMLUExecutorAsync.__init__

def vllm__executor__ray_mlu_executor__RayMLUExecutorAsync____init__(self, *args, **kwargs):
    vllm__executor__ray_mlu_executor__RayMLUExecutorAsync____init____org(self, *args, **kwargs)

    """
    ======================================
    Modified by Chunked Parallel Pipeline.
    ======================================
    For the prefill stage of a request in chunked pipeline parallel, tasks
    in the same pp_rank must be executed in order. Here, we use priority lock
    to implement this function.
    To ensure different requests executed in order, we will reserve a certain
    priority interval for each request. And the interval length is
    `max_model_len`, which is no less than the model execution rounds.
    And for each execution round, the priority is:
        `request_id * max_model_len + model_execution_time`
    """
    self.priority = dict()
    self.priority_interval = self.model_config.max_model_len
    # To ensure pp tasks for the same prefill tokens are created atomically, we
    # use an extra lock to guard it.
    self.lock = asyncio.Lock()
    """
    ======================================
    End by Chunked Parallel Pipeline.
    ======================================
    """

'''
=============================
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__ray_mlu_executor__RayMLUExecutorAsync__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
    ==================
    '''
    assert not self.use_ray_spmd_worker, (
        "RayMLUExecutorAsync is not supported for spmd mode.")
    return await DistributedMLUExecutorAsync.execute_model_async(
        self, 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.
'''
async def vllm__executor__ray_mlu_executor__RayMLUExecutorAsync___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
    ==================
    '''
    assert not self.use_ray_spmd_worker, (
        "driver_worker does not exist for VLLM_USE_RAY_SPMD_WORKER=1")
    if not self.tp_driver_workers:
        return await self.driver_exec_method(
            "execute_model", execute_model_req, prefill_locs, token_chunk_sizes)
    """
    ======================================
    Modified by Chunked Parallel Pipeline.
    ======================================
    Use PriorityLock instead of lock to ensure that tasks in the same pp rank
    are executed with the dispatched order.
    """
    request_id = 'dummy'
    update_priority_threshold = False
    is_prompt = False
    if execute_model_req is not None:
        assert len(execute_model_req.seq_group_metadata_list) == 1, \
            "Only single batch is supported for chunked pipeline parallel mode."
        request_id = execute_model_req.seq_group_metadata_list[0].request_id
        seq_group_metadata = execute_model_req.seq_group_metadata_list[0]
        request_priority = self.priority.setdefault(
            request_id, len(self.priority)*self.model_config.max_model_len)
        seq_data = list(seq_group_metadata.seq_data.values())[0]
        seq_len = seq_data.get_len()

        # Update priority threshold to schedule next request.
        is_prompt = seq_group_metadata.is_prompt
        if is_prompt and seq_len == prefill_locs[0] + token_chunk_sizes[0]:
            update_priority_threshold = True
    else:
        request_priority = -1

    if self.pp_locks is None:
        # This locks each pipeline parallel stage so multiple virtual
        # engines can't execute on the same stage at the same time
        # We create the locks here to avoid creating them in the constructor
        # which uses a different asyncio loop.
        self.pp_locks = [
            PriorityLock(init_priority_threshold=self.model_config.max_model_len,
                         priority_interval=self.priority_interval)
            for _ in range(self.parallel_config.pipeline_parallel_size)
        ]

    async with self.lock:
        tasks = [
            asyncio.create_task(
                _run_task_with_priority_lock(
                    self.driver_exec_method, self.pp_locks[0], request_priority,
                    update_priority_threshold,
                    "execute_model", execute_model_req, prefill_locs, token_chunk_sizes,
                    request_priority))
        ]
        for pp_rank, driver_worker in enumerate(self.tp_driver_workers,
                                                start=1):
            tasks.append(
                asyncio.create_task(
                    _run_task_with_priority_lock(
                        driver_worker.execute_method.remote,
                        self.pp_locks[pp_rank], request_priority,
                        update_priority_threshold,
                        "execute_model", execute_model_req, prefill_locs, token_chunk_sizes,
                        request_priority)))
    if execute_model_req is not None:
        self.priority[request_id] += (token_chunk_sizes[0] if is_prompt else 1)

    """
    ======================================
    End by Chunked Parallel Pipeline.
    ======================================
    """

    results = await asyncio.gather(*tasks)

    # Only the last PP stage has the final results.
    return results[-1]

MluHijackObject.apply_hijack(
    RayMLUExecutorAsync,
    RayMLUExecutorAsync.__init__,
    vllm__executor__ray_mlu_executor__RayMLUExecutorAsync____init__
)
MluHijackObject.apply_hijack(
    RayMLUExecutorAsync,
    RayMLUExecutorAsync.execute_model_async,
    vllm__executor__ray_mlu_executor__RayMLUExecutorAsync__execute_model_async
)
MluHijackObject.apply_hijack(
    RayMLUExecutorAsync,
    RayMLUExecutorAsync._driver_execute_model_async,
    vllm__executor__ray_mlu_executor__RayMLUExecutorAsync___driver_execute_model_async
)
