# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import zmq
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
from vllm.logger import init_logger

from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.request import OmniDiffusionRequest

logger = init_logger(__name__)


class Scheduler:
    _instance = None

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super().__new__(cls)
        return cls._instance

    def initialize(self, od_config: OmniDiffusionConfig):
        if hasattr(self, "context") and not self.context.closed:
            logger.warning("SyncSchedulerClient is already initialized. Re-initializing.")
            self.close()

        self.od_config = od_config
        self.context = zmq.Context()  # Standard synchronous context

        # Initialize MessageQueue for broadcasting requests
        # Assuming all readers are local for now as per current launch_engine implementation
        self.mq = MessageQueue(
            n_reader=od_config.num_gpus,
            n_local_reader=od_config.num_gpus,
            local_reader_ranks=list(range(od_config.num_gpus)),
        )

        self.result_mq = None

    def initialize_result_queue(self, handle):
        # Initialize MessageQueue for receiving results
        # We act as rank 0 reader for this queue
        self.result_mq = MessageQueue.create_from_handle(handle, rank=0)
        logger.info("SyncScheduler initialized result MessageQueue")

    def get_broadcast_handle(self):
        return self.mq.export_handle()

    def add_req(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
        """Sends a request to the scheduler and waits for the response."""
        try:
            # Broadcast request to all workers
            self.mq.enqueue(requests)
            # Wait for result from Rank 0 (or whoever sends it)
            if self.result_mq is None:
                raise RuntimeError("Result queue not initialized")

            output = self.result_mq.dequeue()
            return output
        except zmq.error.Again:
            logger.error("Timeout waiting for response from scheduler.")
            raise TimeoutError("Scheduler did not respond in time.")

    def close(self):
        """Closes the socket and terminates the context."""
        if hasattr(self, "context"):
            self.context.term()


# Singleton instance for easy access
scheduler = Scheduler()