from ldnode.runner.schemas import BaseContext
from ldnode.runner.events import Events


class ResourceController:
    """ 资源控制器 """
    def __init__(self, context: BaseContext, cpu_resource: int, gpu_resource: int):
        self.context = context
        self.cpu_resource = cpu_resource
        self.gpu_resource = gpu_resource

    async def acquire(self, cpu_usage: int, gpu_usage: int) -> bool:
        """ 获取资源 """
        if self.cpu_resource >= cpu_usage and self.gpu_resource >= gpu_usage:
            self.cpu_resource -= cpu_usage
            self.gpu_resource -= gpu_usage
            await self.context.submit_event("calculator", Events.LOCK_RESOURCE, {
                "cpu_resource": cpu_usage,
                "gpu_resource": gpu_usage,
            })
            return True
        else:
            return False
        
    async def release(self, cpu_usage: int, gpu_usage: int) -> None:
        """ 释放资源 """
        self.cpu_resource += cpu_usage
        self.gpu_resource += gpu_usage
        await self.context.submit_event("calculator", Events.RELEASE_RESOURCE, {
            "cpu_resource": cpu_usage,
            "gpu_resource": gpu_usage,
        })
