import asyncio
from typing import (List, Optional, Union)

from vllm.envs import VLLM_ENGINE_ITERATION_TIMEOUT_S as ENGINE_ITERATION_TIMEOUT_S
from vllm.core.scheduler import ScheduledSequenceGroup
from vllm.engine.async_timeout import asyncio_timeout
from vllm.outputs import EmbeddingRequestOutput, RequestOutput
from vllm.sequence import ExecuteModelRequest, SequenceGroup, SequenceGroupMetadata
from vllm.engine.async_llm_engine import (_AsyncLLMEngine, AsyncLLMEngine)
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.engine.llm_engine import LLMEngine

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

logger = init_logger(__name__)


def vllm__engine__async_llm_engine___AsyncLLMEngine____init__(self, *args, **kwargs):
    LLMEngine.__init__(self, *args, **kwargs)

    """
    =============================
    Modifies by vllm_mlu
    =============================
    @brief:  Add a member variable to record parallel chunked prefill tasks,
    in which each member means (virtual_engine -> {req_id: task_list})
    """
    self.step_tasks = [dict() for _ in range(len(self.scheduler))]
    """
    =============================
    End of MLU Hijack
    =============================
    """

def _update_scheduler_status(
    self,
    scheduled_seq_groups: List[ScheduledSequenceGroup],
    ignored_seq_groups: List[SequenceGroup],
    seq_group_metadata_list: List[SequenceGroupMetadata]
) -> None:
    """Update scheduler status after emitting prefill task.

    For chunked pipeline parallel, since chunked prefill tasks
    are executed asynchronously, we update scheduler status once
    tasks are emited.
    """
    # Update the scheduled sequence groups.
    for scheduled_seq_group, seq_group_meta in zip(
            scheduled_seq_groups, seq_group_metadata_list):
        seq_group = scheduled_seq_group.seq_group
        seq_group.update_num_computed_tokens(
            scheduled_seq_group.token_chunk_size)

    # Free the finished sequence groups.
    for scheduler in self.scheduler:
        scheduler.free_finished_seq_groups()

async def vllm__engine__async_llm_engine___AsyncLLMEngine__step_async(
    self, virtual_engine: int
) -> Optional[List[Union[RequestOutput, EmbeddingRequestOutput]]]:
    """Performs one decoding iteration and returns newly generated results.
    The workers are ran asynchronously if possible.

    This function performs one decoding iteration of the engine. It first
    schedules the sequences to be executed in the next iteration and the
    token blocks to be swapped in/out/copy. Then, it executes the model
    and updates the scheduler with the model outputs. Finally, it decodes
    the sequences and returns the newly generated results.
    """
    # these are cached outputs from previous iterations. None if on first
    # iteration
    cached_outputs = self.cached_scheduler_outputs[virtual_engine]
    seq_group_metadata_list = cached_outputs.seq_group_metadata_list
    scheduler_outputs = cached_outputs.scheduler_outputs
    allow_async_output_proc = cached_outputs.allow_async_output_proc

    ctx = self.scheduler_contexts[virtual_engine]

    # Clear outputs for each new scheduler iteration
    ctx.request_outputs.clear()

    # skip the scheduler if there are any remaining steps in the seq groups.
    # This ensures that the scheduler is only called again when the current
    # batch has completed.
    if not self._has_remaining_steps(seq_group_metadata_list):

        # Schedule iteration
        (seq_group_metadata_list, scheduler_outputs,
         allow_async_output_proc
         ) = self.scheduler[virtual_engine].schedule()

        ctx.seq_group_metadata_list = seq_group_metadata_list
        ctx.scheduler_outputs = scheduler_outputs

        # Maybe switch from async mode to sync mode
        if not allow_async_output_proc and len(ctx.output_queue) > 0:
            self._process_model_outputs(ctx=ctx)

        if (self.scheduler_config.is_multi_step
                and scheduler_outputs.num_lookahead_slots > 0):
            # cache the scheduler outputs for the next iteration if we have
            # lookahead slots
            self._cache_scheduler_outputs_for_multi_step(
                virtual_engine, seq_group_metadata_list, scheduler_outputs,
                allow_async_output_proc)

    assert seq_group_metadata_list is not None
    assert scheduler_outputs is not None

    if not scheduler_outputs.is_empty():
        finished_requests_ids = self.scheduler[
            virtual_engine].get_and_reset_finished_requests_ids()

        # Check if we have a cached last_output from the previous iteration.
        # For supporting PP this is probably the best way to pass the
        # sampled_token_ids, as a separate broadcast over all the PP stages
        # will cause one virtual engine's microbatch to block the pipeline.
        last_sampled_token_ids = \
            self._get_last_sampled_token_ids(virtual_engine)

        execute_model_req = ExecuteModelRequest(
            seq_group_metadata_list=seq_group_metadata_list,
            blocks_to_swap_in=scheduler_outputs.blocks_to_swap_in,
            blocks_to_swap_out=scheduler_outputs.blocks_to_swap_out,
            blocks_to_copy=scheduler_outputs.blocks_to_copy,
            virtual_engine=virtual_engine,
            num_lookahead_slots=scheduler_outputs.num_lookahead_slots,
            running_queue_size=scheduler_outputs.running_queue_size,
            finished_requests_ids=finished_requests_ids,
            # We use ExecuteModelRequest to pass the last sampled_token_ids
            # to each of the non-last PP stages for in-place prepare_input.
            last_sampled_token_ids=last_sampled_token_ids)

        if allow_async_output_proc:
            execute_model_req.async_callback = self.async_callbacks[
                virtual_engine]

        # Execute the model.
        """
        =============================
        Modifies by vllm_mlu
        =============================
        @brief: for chunked prefill tasks except the final task for a single
        request, create them asynchronously. And for the last prefill task,
        gather all previous tasks and get the final output.
        """
        if seq_group_metadata_list[0].is_prompt:
            assert len(seq_group_metadata_list) == 1, \
                "Currently we only support schedule single batch in " \
                "prefill stage for chunked pipeline parallel."
            token_chunk_size = seq_group_metadata_list[0].token_chunk_size
            seq_data = list(seq_group_metadata_list[0].seq_data.values())[0]
            prefill_loc = seq_data.get_num_computed_tokens()
            task = asyncio.create_task(
                self.model_executor.execute_model_async(execute_model_req, [prefill_loc], [token_chunk_size]))
            request_id = seq_group_metadata_list[0].request_id
            self.step_tasks[virtual_engine].setdefault(request_id, []).append(task)

            # Gather point: if all prefill tasks for current sequence group
            # have been dispatched, we wait all prompt tasks and get the
            # final output.
            seq_len = seq_data.get_len()
            if token_chunk_size + prefill_loc == seq_len:
                outputs = await asyncio.gather(*self.step_tasks[virtual_engine][request_id])
                outputs = outputs[-1]
            else:
                # Since prefill stage has not been completely finished, we
                # just update scheduler and sequence status and return None.
                _update_scheduler_status(self, scheduler_outputs.scheduled_seq_groups,
                    scheduler_outputs.ignored_seq_groups, seq_group_metadata_list)
                return None
        else:
            """
            =============================
            End of MLU Hijack
            =============================
            """
            outputs = await self.model_executor.execute_model_async(
                execute_model_req)

        # we need to do this here so that last step's sampled_token_ids can
        # be passed to the next iteration for PP.
        if self.scheduler_config.is_multi_step:
            self._update_cached_scheduler_output(virtual_engine, outputs)
    else:
        if len(ctx.output_queue) > 0:
            self._process_model_outputs(ctx=ctx)
        outputs = []

    # Finish the current step for all the sequence groups.
    if self.scheduler_config.is_multi_step:
        for seq_group in seq_group_metadata_list:
            seq_group.finish_step()

    if not self._has_remaining_steps(seq_group_metadata_list):
        # Clear the cache if we have finished all the steps
        if self.scheduler_config.is_multi_step:
            self.cached_scheduler_outputs[
                virtual_engine] = SchedulerOutputState()

        # is_first_step_output is True only when the num_steps of all
        # the sequences are 1. When the num_steps > 1,
        # multi_step_model_runner does the first-step output append.
        is_first_step_output: bool = False if not seq_group_metadata_list \
            else seq_group_metadata_list[0].state.num_steps == 1

        ctx.append_output(outputs=outputs,
                          seq_group_metadata_list=seq_group_metadata_list,
                          scheduler_outputs=scheduler_outputs,
                          is_async=allow_async_output_proc,
                          is_last_step=True,
                          is_first_step_output=is_first_step_output)

        if outputs and allow_async_output_proc:
            assert len(
                outputs
            ) == 1, "Async postprocessor expects only a single output set"
            self._advance_to_next_step(
                outputs[0], seq_group_metadata_list,
                scheduler_outputs.scheduled_seq_groups)

        if not allow_async_output_proc:
            self._process_model_outputs(ctx=ctx)

            # Log stats.
            self.do_log_stats(scheduler_outputs, outputs)

            # Tracing
            self.do_tracing(scheduler_outputs)

    else:
        # Multi-step case
        return ctx.request_outputs

    if not self.has_unfinished_requests():
        # Drain async postprocessor (if exists)
        if len(ctx.output_queue) > 0:
            self._process_model_outputs(ctx=ctx)
        assert len(ctx.output_queue) == 0

    return ctx.request_outputs

async def vllm__engine__async_llm_engine__AsyncLLMEngine__engine_step(
    self, virtual_engine: int
) -> bool:
    """Kick the engine to process the waiting requests.

    Returns True if there are in-progress requests."""

    new_requests, aborted_requests = (
        self._request_tracker.get_new_and_aborted_requests())

    for new_request in new_requests:
        # Add the request into the vLLM engine's waiting queue.
        try:
            await self.engine.add_request_async(**new_request)
        except ValueError as e:
            # TODO: use a vLLM specific error for failed validation
            self._request_tracker.process_exception(
                new_request["request_id"],
                e,
                verbose=self.log_requests,
            )

    if aborted_requests:
        await self._engine_abort(aborted_requests)

    request_outputs = await self.engine.step_async(virtual_engine)

    """
    ======================================
    Modified by Chunked Parallel Pipeline.
    ======================================
    When request_outputs is None, it means prefill tasks are not finished.
    """
    if request_outputs is None:
        return True
    """
    ======================================
    End by Chunked Parallel Pipeline.
    ======================================
    """
    # Put the outputs into the corresponding streams.
    # If used as a callback, then already invoked inside
    # LLMEngine's _process_model_outputs
    if not self.use_process_request_outputs_callback:
        all_finished = self.process_request_outputs(request_outputs)
    else:
        # For callback case, we only need to detect when all
        # requests are finished
        all_finished = all(request_output.finished
                           for request_output in request_outputs)

    return not all_finished

MluHijackObject.apply_hijack(
    _AsyncLLMEngine,
    _AsyncLLMEngine.__init__,
    vllm__engine__async_llm_engine___AsyncLLMEngine____init__
)
MluHijackObject.apply_hijack(
    _AsyncLLMEngine,
    _AsyncLLMEngine.step_async,
    vllm__engine__async_llm_engine___AsyncLLMEngine__step_async
)
MluHijackObject.apply_hijack(
    AsyncLLMEngine,
    AsyncLLMEngine.engine_step,
    vllm__engine__async_llm_engine__AsyncLLMEngine__engine_step
)
