import asyncio

from ldnode import logger
from ldnode.runner.config import Settings
from ldnode.runner.schemas import Service, BaseContext
from ldnode.runner.task_manager import TaskManager
from ldnode.runner.events import (
    on, 
    Events, 
    TaskSubmitted,
    TaskFinished,
    TaskCanceled,
    TaskCancel,
    WorkerStateChanged,
)
from .application import (
    CPUScheduleController,
    ResourceController,
    LoadCounter,
)
from .handlers import handler_registry
from .worker import CalculatorWorker


class CalculatorService(Service):
    """ 计算服务 """
    def __init__(self, settings: Settings, context: BaseContext):
        super().__init__(settings, context)
        self.running = False
        self.tasks = TaskManager()
        self.exit_flag = asyncio.Event()
        self.waiting_tasks: list[TaskSubmitted] = []
        self.schedule_controller = CPUScheduleController(
            self.context, 
            cpu_resource=self.settings.CPU_RESOURCE, 
            executor_cls=CalculatorWorker, 
            running_checker=lambda: self.running,
        )
        self.resource_controller = ResourceController(
            self.context,
            cpu_resource=self.settings.CPU_RESOURCE, 
            gpu_resource=0,
        )
        self.load_counter = LoadCounter(
            self.context,
            load_querier=lambda: self.schedule_controller.load,
            running_checker=lambda: self.running,
        )

    async def start(self) -> None:
        """ 启动服务 """
        self.running = True
        # 启动调度器
        await self.schedule_controller.schedule()
        # 启动负载计算
        await self.load_counter.count()
        # 等待退出
        await self.exit_flag.wait()
    
    async def shutdown(self) -> None:
        """ 关闭服务 """
        self.running = False
        self.exit_flag.set()
        
    async def release_next_task(self) -> None:
        """ 释放下一个等待中的任务 """
        if self.waiting_tasks:
            next_task = self.waiting_tasks.pop(0)
            self.tasks.create_task(self.handle_task_submitted("calculator", next_task))
        
    @on(Events.WORKER_START)
    async def handle_worker_start(self, _: str, data: WorkerStateChanged) -> None:
        """ 处理工作进程启动 """
        await self.schedule_controller.awake_executor(data['worker_id'])

    @on(Events.WORKER_EXIT)
    async def handle_worker_exit(self, _: str, data: WorkerStateChanged) -> None:
        """ 处理工作进程退出 """
        await self.schedule_controller.sleep_executor(data['worker_id'])
        
    @on(Events.TASK_SUBMITTED)
    async def handle_task_submitted(self, _: str, task: TaskSubmitted) -> None:
        """ 处理任务提交 """
        if handler := handler_registry.get(task['method']):
            if await self.resource_controller.acquire(handler.cpu_usage, handler.gpu_usage):
                # 获取资源成功, 提交任务
                await self.schedule_controller.submit(task)
            else:
                # 进入队列
                self.waiting_tasks.append(task)
        else:
            logger.calculator.warning(f"未找到处理器: {task['method']}")

    @on(Events.TASK_FINISHED)
    async def handle_task_finished(self, source: str, data: TaskFinished) -> None:
        """ 处理任务完成 """
        logger.calculator.info(f"{source} 完成任务: {data['method']}({data['task_id']})")
        if handler := handler_registry.get(data['method']):
            # 释放资源
            await self.resource_controller.release(handler.cpu_usage, handler.gpu_usage)
            # 释放下一个等待中的任务
            await self.release_next_task()
            
    @on(Events.TASK_CANCEL)
    async def handle_task_cancel(self, _: str, data: TaskCancel) -> None:
        """ 处理任务取消 """
        logger.calculator.info(f"取消任务: {data['method']}({data['task_id']})")
        # 判断任务是否还在等待队列中
        if task := next(filter(lambda task: task['task_id'] == data['task_id'], self.waiting_tasks), None):
            # 取消任务
            self.waiting_tasks.remove(task)
            await self.context.submit_event('calculator', Events.TASK_CANCELED, {
                "task_id": data['task_id'],
                "method": data['method'],
                "is_released": False,
            })
        else:
            # 任务不在等待队列中
            await self.schedule_controller.cancel_task(data['task_id'])
            
    @on(Events.TASK_CANCELED)
    async def handle_task_canceled(self, _: str, data: TaskCanceled) -> None:
        """ 处理任务取消 """
        logger.calculator.info(f"任务取消: {data['method']}({data['task_id']})")
        if data['is_released']:
            if handler := handler_registry.get(data['method']):
            # 释放资源
                await self.resource_controller.release(handler.cpu_usage, handler.gpu_usage)
            else:
                logger.calculator.warning(f"未找到处理器: {data['method']}")
