import threading as th
import queue
from uuid import uuid4
from typing import Callable, Optional, Any
from dataclasses import dataclass

from ldnode.runner.events import TaskSubmitted
from ldnode.calculator.schemas import CancellationToken as AbstractCancellationToken
from .models import (
    OutputMessageType,
    OutputMessage,
)


class CancellationToken(AbstractCancellationToken):
    """ 取消令牌 """
    def __init__(self):
        self._canceled = False
        
    def cancel(self) -> None:
        """ 取消 """
        self._canceled = True
    
    def is_canceled(self) -> bool:
        """ 是否取消 """
        return self._canceled


@dataclass
class PendingInvoke:
    """ 调用等待对象 """
    done: th.Event # 完成事件
    result: Any = None # 结果


class WorkerContext:
    """ 工作进程上下文 """
    def __init__(self, running_checker: Callable[[], bool], exit_handler: Callable[[], None]):
        # 外部接口
        self.check_running = running_checker
        self.handle_exit = exit_handler
        # 任务调度
        self.current_task: Optional[TaskSubmitted] = None
        self.current_cancel_token: Optional[CancellationToken] = None
        self.calc_signal = th.Event()
        # 输出信箱
        self.outbox: "queue.Queue[OutputMessage]" = queue.Queue()
        # 计算调用
        self.pending_invokes: dict[str, PendingInvoke] = {}
        self.pending_lock = th.Lock()
        
    @property
    def running(self) -> bool:
        """ 运行状态 """
        return self.check_running()
    
    def exit(self) -> None:
        """ 退出 """
        self.handle_exit()
    
    def load_task(self, task: TaskSubmitted) -> None:
        """ 加载任务 """
        self.current_task = task
        self.current_cancel_token = CancellationToken()
        self.calc_signal.set()
        
    def wait_for_task(self) -> tuple[TaskSubmitted, CancellationToken]:
        """ 等待任务 """
        self.calc_signal.wait()
        self.calc_signal.clear()
        return self.current_task, self.current_cancel_token
    
    def cancel_task(self) -> None:
        """ 取消任务 """
        if self.current_cancel_token:
            self.current_cancel_token.cancel()
        
    def is_canceled(self) -> bool:
        """ 是否取消 """
        return self.current_cancel_token and self.current_cancel_token.is_canceled()
        
    def submit_invoke_answer(self, invoke_id: str, result: Any) -> None:
        """ 提交调用应答 """
        with self.pending_lock:
            if pending_invoke := self.pending_invokes.pop(invoke_id, None):
                pending_invoke.result = result
                pending_invoke.done.set()
    
    def invoke(self, method: str, parameters: dict[str, Any], timeout: float) -> Any:
        """ 调用计算方法 """
        # 创建调用
        pending_invoke = PendingInvoke(done=th.Event())
        invoke_id = str(uuid4())
        with self.pending_lock:
            self.pending_invokes[invoke_id] = pending_invoke
        # 发送调用
        self.outbox.put({
            "msg_type": OutputMessageType.INVOKE,
            "payload": {
                "task_id": invoke_id,
                "method": method,
                "parameters": parameters,
            },
        })
        # 等待调用结果
        ok = pending_invoke.done.wait(timeout)
        if not ok:
            # 超时时清理挂起记录，避免内存泄漏和孤儿回包问题
            with self.pending_lock:
                self.pending_invokes.pop(invoke_id, None)
            return {"data_type": "error", "data": "调用超时"}
        return pending_invoke.result
