from .receive_scheduler import ReceiveScheduler
from .task_scheduler import TaskScheduler
from ..model import MessageItem
from typing import Callable
from ..queue_operation import QueueOperation
from .cleanup_scheduler import CleanupScheduler

class QueueScheduler:
    def __init__(self, receive_thread_num=1, task_thread_num=1, shard_num=4):
        self.queue_operation = QueueOperation(shard_num)
        self.receive_scheduler = ReceiveScheduler(self.queue_operation, receive_thread_num)
        self.task_scheduler = TaskScheduler(self.queue_operation, task_thread_num)
        self.cleanup_scheduler = CleanupScheduler(self.queue_operation)

    def send_message(self, message: MessageItem, callback: Callable):
        self.receive_scheduler.send_message(message, callback)

    def start_queue_scheduler(self):
        self.receive_scheduler.start_receive_thread()
        self.task_scheduler.start_task_thread()
        self.cleanup_scheduler.start_cleanup()

    def stop_queue_scheduler(self):
        self.receive_scheduler.stop_receive_thread()
        self.task_scheduler.stop_task_thread()
        self.cleanup_scheduler.stop_cleanup()
        # self.queue_operation.close_all_connections()

__all__ = [
    "QueueScheduler"
]