from .receive_scheduler import ReceiveScheduler
from .task_scheduler import TaskScheduler
from .listen_data_scheduler import ListenDataScheduler
from ...model import MessageItem
from typing import Callable
from ...queue_operation import QueueOperation
from ...queue_operation.listen_operation import ListenOperation
from ..cleanup_scheduler import CleanupScheduler
import os
from ..base import BaseScheduler


class StandardQueueScheduler(BaseScheduler):
    def __init__(self, receive_thread_num=1, task_thread_num=1, shard_num=4):
        self.queue_operation = QueueOperation(shard_num)
        self.listen_operation = ListenOperation(
            os.path.join(self.queue_operation.db_dir, "listen.db")
        )
        self.listen_operation.create_table()

        self.listen_scheduler = ListenDataScheduler(self.listen_operation)
        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 update_listen_data(self, key, value):
        self.listen_operation.update_listen_data(key, value)

    def get_listen_datas(self):
        return self.listen_operation.get_values()

    def get_listen_data(self, key):
        return self.listen_operation.get_value(key)

    def start(self):
        self.receive_scheduler.start_receive_thread()
        self.task_scheduler.start_task_thread()
        self.cleanup_scheduler.start_cleanup()
        self.listen_scheduler.start_listen_data()

    def stop(self):
        self.receive_scheduler.stop_receive_thread()
        self.listen_scheduler.stop_listen_data()
        self.task_scheduler.stop_task_thread()
        self.cleanup_scheduler.stop_cleanup()
        # self.queue_operation.close_all_connections()


__all__ = ["StandardQueueScheduler"]
