import pickle
from typing import Iterator, List, Optional, Union

import cloudpickle
import zmq

from vllm import SamplingParams
# yapf conflicts with isort for this block
# yapf: disable
from vllm.engine.multiprocessing import (RPCAbortRequest, RPCProcessRequest,
                                         RPCUProfileRequest)
from vllm.engine.llm_engine import LLMEngine
# yapf conflicts with isort for this block
# yapf: disable
from vllm.engine.multiprocessing import (IPC_DATA_EXT, IPC_HEALTH_EXT,
                                         IPC_INPUT_EXT, IPC_OUTPUT_EXT,
                                         RPCAbortRequest, RPCProcessRequest,
                                         RPCUProfileRequest)
from vllm.engine.multiprocessing.engine import (MQLLMEngine,
                                                POLLING_TIMEOUT_MS)
from vllm.logger import init_logger

from vllm_mlu.engine.multiprocessing import RPCSchedulerProfileRequest
from vllm_mlu.mlu_hijack_utils import MluHijackObject

logger = init_logger(__name__)


vllm__engine__multiprocessing__engine__MQLLMEngine____init____org = MQLLMEngine.__init__

class MQLLMEngine_V2(MQLLMEngine):

    def __init__(self,
                 ipc_path: str,
                 use_async_sockets: bool,
                 *args,
                 log_requests: bool = True,
                 **kwargs) -> None:
        # For MQLLMEngine, we can use cached outputs, since each new request
        # output is immediately pickled and send over the socket, which frees
        # the python object to be reused again.
        kwargs['use_cached_outputs'] = True

        self.engine = LLMEngine(*args, **kwargs)
        self.log_requests = log_requests

        self.use_async_sockets = use_async_sockets
        if self.use_async_sockets:
            self.engine.process_request_outputs_callback = \
                self._async_socket_engine_callback

        self.ctx = zmq.Context()  # type: ignore[attr-defined]

        # Receive input from the client.
        self.input_socket = self.ctx.socket(zmq.constants.PULL)
        self.input_socket.bind(f"{ipc_path}{IPC_INPUT_EXT}")

        # Send output stream back to client.
        self.output_socket = self.ctx.socket(zmq.constants.PUSH)
        self.output_socket.bind(f"{ipc_path}{IPC_OUTPUT_EXT}")

        # Send heartbeats back to client.
        self.heartbeat_socket = self.ctx.socket(zmq.constants.PUSH)
        self.heartbeat_socket.bind(f"{ipc_path}{IPC_HEALTH_EXT}")

        # IPC path for the data socket.
        self.data_ipc_path = f"{ipc_path}{IPC_DATA_EXT}"

        # Error state.
        self._errored_with: Optional[BaseException] = None

        self.collect_scheduler_view = False

    def run_engine_loop(self):
        """Core busy loop of the LLMEngine."""

        while True:
            if not self.engine.has_unfinished_requests():
                # Poll until there is work to do.
                while self.input_socket.poll(timeout=POLLING_TIMEOUT_MS) == 0:
                    # When there's no work, check on engine health and send
                    # health status back to client
                    self._health_check()
                    self.engine.do_log_stats()
                    logger.debug("Waiting for new requests in engine loop.")

            # Handle any input from the client.
            self.handle_new_input()

            '''
            =============================
            Add by vllm_mlu
            =============================
            @brief: support scheduler view
            '''
            if self.collect_scheduler_view:
                self.collect_scheduler_view = False
                continue
            '''
            ==================
            End of MLU Hijack
            ==================
            '''

            # Engine step.
            request_outputs = self.engine_step()

            # Send request outputs (if async, done in engine_step callback).
            if not self.use_async_sockets:
                self._send_outputs(request_outputs)

    def handle_new_input(self):
        """Handle new input from the socket"""
        try:
            while self.input_socket.poll(timeout=0) != 0:
                frames = self.input_socket.recv_multipart(copy=False)
                request = pickle.loads(frames[0].buffer)

                '''
                =============================
                Add by vllm_mlu
                =============================
                @brief: support scheduler view
                '''

                if isinstance(request, RPCProcessRequest):
                    if len(frames) > 1:
                        # Use cloudpickle for logits processors
                        assert isinstance(request.params, SamplingParams)
                        lprocs = cloudpickle.loads(frames[1].buffer)
                        request.params.logits_processors = lprocs
                    self._handle_process_request(request)
                elif isinstance(request, RPCAbortRequest):
                    self._handle_abort_request(request)
                elif isinstance(request, RPCUProfileRequest):
                    if request == RPCUProfileRequest.START_PROFILE:
                        self.start_profile()
                    else:
                        self.stop_profile()
                elif isinstance(request, RPCSchedulerProfileRequest):
                    self.collect_scheduler_view = True
                    if request == RPCSchedulerProfileRequest.INIT_SCHEDULER_VIEW:
                        self.init_scheduler_view()
                    elif request == RPCSchedulerProfileRequest.SAVE_SCHEDULER_VIEW:
                        self.save_scheduler_view()
                else:
                    raise ValueError("Unknown RPCRequest Type: "
                                     f"{type(request)}")
                '''
                ==================
                End of MLU Hijack
                ==================
                '''

        except Exception as e:
            self._set_errored(e)
            self._send_unhealthy(e)
            raise e

    def init_scheduler_view(self):
        """Init scheduler view."""
        self.engine.init_scheduler_view()

    def save_scheduler_view(self):
        """Save scheduler view."""
        self.engine.save_scheduler_view()


MluHijackObject.apply_hijack(MQLLMEngine,
                             MQLLMEngine.__init__,
                             MQLLMEngine_V2.__init__)
MluHijackObject.apply_hijack(MQLLMEngine,
                             MQLLMEngine.run_engine_loop,
                             MQLLMEngine_V2.run_engine_loop)
MluHijackObject.apply_hijack(MQLLMEngine,
                             MQLLMEngine.handle_new_input,
                             MQLLMEngine_V2.handle_new_input)
MluHijackObject.apply_hijack(MQLLMEngine,
                             "init_scheduler_view",
                             MQLLMEngine_V2.init_scheduler_view)
MluHijackObject.apply_hijack(MQLLMEngine,
                             "save_scheduler_view",
                             MQLLMEngine_V2.save_scheduler_view)