import threading as th
import multiprocessing as mp
import queue
from typing import Optional

from ldnode import logger
from ldnode.runner.events import TaskSubmitted
from .models import (
    InputMessageType, 
    InputMessage, 
    InvokeResult,
)
from .worker_context import WorkerContext


class RecvController(th.Thread):
    """ 接收控制器 """
    def __init__(self, context: WorkerContext, input_channel: mp.Queue):
        super().__init__(daemon=True)
        self.context = context
        self.input_channel = input_channel
        
    def exit_if_main_notify(self, msg: Optional[InputMessage]) -> bool:
        """ 如果收到主进程通知，则退出 """
        if msg is None:
            self.context.exit()
            return True
        else:
            return False
        
    def run(self) -> None:
        """ 运行控制器 """
        while self.context.running:
            try:
                msg: Optional[InputMessage] = self.input_channel.get(block=True, timeout=1)
                if self.exit_if_main_notify(msg):
                    break
                match msg['msg_type']:
                    case InputMessageType.TASK_SUBMITTED:
                        task_submitted: TaskSubmitted = msg['payload']
                        self.context.load_task(task_submitted)
                    case InputMessageType.TASK_CANCELLED:
                        self.context.cancel_task()
                    case InputMessageType.INVOKE_ANSWER:
                        result: InvokeResult = msg['payload']
                        self.context.submit_invoke_answer(result['task_id'], result['result'])
            except queue.Empty:
                continue
            except Exception as e:
                logger.calculator.error(f"接收消息失败: {e}")
