import time
import queue
from typing import Type, Callable, Optional

from ldnode import logger
from ldnode.runner.schemas import (
    BaseContext,
    BaseWorkerState,
    Worker,
)
from ldnode.runner.task_manager import TaskManager
from ldnode.runner.events import (
    Events,
    TaskSubmitted,
    TaskAnswer,
)
from .models import (
    InputMessageType,
    OutputMessageType,
    OutputMessage,
    InvokeResult,
)


class CPUScheduleController:
    """ CPU调度控制器 """
    def __init__(self, context: BaseContext, cpu_resource: int, executor_cls: Type[Worker], running_checker: Callable[[], bool]):
        self.context = context
        self.cpu_resource = cpu_resource
        self.executor_cls = executor_cls
        self.running_checker = running_checker
        # 记录执行器数量
        self.worker_count = 0
        # 记录执行器
        self.executors: dict[str, BaseWorkerState] = {}
        # 记录空闲的执行器
        self.free_executors: set[str] = set()
        # 记录正在忙碌的执行器
        self.busy_executors: dict[str, TaskSubmitted] = {}
        # 记录正在等待执行器的任务
        self.waiting_queue: list[TaskSubmitted] = []
        # 记录正在等待的调用和执行器的关系
        self.pending_invokes: dict[str, str] = {}
        # 异步任务管理器
        self.tasks = TaskManager()
        
    @property
    def load(self) -> int:
        """ 负载 """
        return len(self.busy_executors) + len(self.waiting_queue)
        
    async def _submit_task(self, executor_id: str, task: TaskSubmitted) -> None:
        """ 提交任务到执行器 """
        if executor_id in self.executors:
            await self.executors[executor_id].send({
                "msg_type": InputMessageType.TASK_SUBMITTED,
                "payload": task,
            })
            
    async def _cancel_task(self, executor_id: str) -> None:
        """ 取消任务到执行器 """
        if executor_id in self.executors:
            await self.executors[executor_id].send({
                "msg_type": InputMessageType.TASK_CANCELLED,
                "payload": None,
            })
        
    async def _submit_invoke_answer(self, executor_id: str, result: InvokeResult) -> None:
        """ 提交调用应答到执行器 """
        if executor_id in self.executors:
            await self.executors[executor_id].send({
                "msg_type": InputMessageType.INVOKE_ANSWER,
                "payload": result,
            })
        
    async def _handle_task_answer(self, executor_id: str, task_answer: TaskAnswer) -> None:
        """ 处理任务应答 """
        if task_answer['task_id'] in self.pending_invokes:
            # 提交内部调用应答
            source_executor_id = self.pending_invokes[task_answer['task_id']]
            await self._submit_invoke_answer(source_executor_id, {
                "task_id": task_answer['task_id'],
                "result": task_answer['result'],
            })
            if task_answer['completed']:
                self.pending_invokes.pop(task_answer['task_id'], None)
        else:
            # 提交外部调用应答
            await self.context.submit_event(executor_id, Events.TASK_ANSWER, task_answer)
            # 如果任务完成 通知调度器任务完成
            if task_answer['completed']:
                if task_answer.get('result', {}).get('canceled', False):
                    await self.context.submit_event(executor_id, Events.TASK_CANCELED, {
                        "task_id": task_answer['task_id'],
                        "method": task_answer['method'],
                        "is_released": False,
                    })
                await self.context.submit_event(executor_id, Events.TASK_FINISHED, {
                    "task_id": task_answer['task_id'],
                    "method": task_answer['method'],
                })
        # 如果任务完成 并且执行器正在忙碌 则释放执行器
        if task_answer['completed'] and executor_id in self.busy_executors:
            self.busy_executors.pop(executor_id)
            self.free_executors.add(executor_id)
            # 释放下一个等待中的任务
            await self._release_waiting_task()
                
            
    async def _handle_invoke(self, executor_id: str, task: TaskSubmitted) -> None:
        """ 处理调用 """
        self.pending_invokes[task['task_id']] = executor_id
        await self.submit(task)
        
    async def _accept_executor(self, executor: BaseWorkerState) -> None:
        """ 接受CPU执行器 """
        executor_id = executor.get_worker_id()
        while self.running_checker() and executor.is_used():
            try:
                output: Optional[OutputMessage] = await executor.recv()
                if output is not None:
                    match output['msg_type']:
                        case OutputMessageType.TASK_ANSWER:
                            await self._handle_task_answer(executor_id, output['payload'])
                        case OutputMessageType.INVOKE:
                            await self._handle_invoke(executor_id, output['payload'])
                        case OutputMessageType.EVENT:
                            pass
            except queue.Empty:
                continue
            except Exception as e:
                logger.calculator.warning(f"{executor.get_worker_id()} 监听执行器输出失败: {e}")
        
    async def _spawn_executor(self) -> None:
        """ 启动CPU执行器 """
        worker_id = ""
        try:
            worker_index = self.worker_count
            worker_state = await self.context.register_worker("calculator", worker_index, self.executor_cls)
            worker_id = worker_state.get_worker_id()
            self.executors[worker_id] = worker_state
            self.tasks.create_task(self._accept_executor(worker_state))
            self.worker_count += 1
            self.latest_worker_start_at = time.time()
        except Exception as e:
            logger.calculator.warning(f"创建计算进程失败: {e}")
            if worker_id:
                self.executors.pop(worker_id, None)
                
    async def _terminate_executor(self, worker_id: str) -> None:
        """ 终止CPU执行器 """
        if worker_id in self.executors:
            await self.executors[worker_id].close()
            del self.executors[worker_id]
            
    async def _release_waiting_task(self) -> None:
        """ 释放等待中的任务 """
        if self.waiting_queue:
            task = self.waiting_queue.pop(0)
            await self.submit(task)
        
    async def schedule(self) -> None:
        """ 开始调度 """
        # 为CPU资源创建对应的执行器
        for _ in range(self.cpu_resource):
            await self._spawn_executor()
            
    async def submit(self, task: TaskSubmitted) -> None:
        """ 提交任务 """
        if self.free_executors:
            # 获取空闲执行器
            free_executor_id = self.free_executors.pop()
            # 将执行器和任务放入忙碌队列
            self.busy_executors[free_executor_id] = task
            # 提交任务到执行器
            await self._submit_task(free_executor_id, task)
        else:
            # 将任务加入等待队列
            self.waiting_queue.append(task)
            
    async def awake_executor(self, executor_id: str) -> None:
        """ 唤醒执行器 """
        if executor_id in self.executors and executor_id not in self.free_executors:
            # 将执行器加入空闲队列
            self.free_executors.add(executor_id)
            # 释放等待中的任务
            await self._release_waiting_task()
             
    async def sleep_executor(self, executor_id: str) -> None:
        """ 睡眠执行器 """
        if executor_id in self.executors:
            # 当执行器处于空闲状态 则立即将它移除
            if executor_id in self.free_executors:
                self.free_executors.remove(executor_id)
            # 当执行器处于忙碌状态 则将它从忙碌状态中移除 并执行器正在运行任务重新放入等待队列
            if executor_id in self.busy_executors:
                task = self.busy_executors.pop(executor_id)
                self.waiting_queue.append(task)

    async def cancel_task(self, task_id: str) -> None:
        """ 取消任务 """
        # 判断任务是否还在等待队列中
        if task := next(filter(lambda task: task['task_id'] == task_id, self.waiting_queue), None):
            # 取消任务
            self.waiting_queue.remove(task)
            await self.context.submit_event('calculator', Events.TASK_CANCELED, {
                "task_id": task_id,
                "method": task['method'],
                "is_released": True,
            })
        else:
            # 找到正在执行的执行器
            if executor_id := next(filter(
                lambda executor_id: self.busy_executors[executor_id]['task_id'] == task_id, 
                self.busy_executors.keys()), None):
                await self._cancel_task(executor_id)
