from ldnode import logger
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:
            before_cpu = self.cpu_resource
            before_gpu = self.gpu_resource
            self.cpu_resource -= cpu_usage
            self.gpu_resource -= gpu_usage
            logger.calculator.info(
                "[ResourceLocal] 锁定本地资源 "
                f"cpu_usage={cpu_usage} gpu_usage={gpu_usage} "
                f"before_cpu={before_cpu} after_cpu={self.cpu_resource} "
                f"before_gpu={before_gpu} after_gpu={self.gpu_resource}"
            )
            await self.context.submit_event("calculator", Events.LOCK_RESOURCE, {
                "cpu_resource": cpu_usage,
                "gpu_resource": gpu_usage,
            })
            return True
        else:
            logger.calculator.warning(
                "[ResourceLocal] 本地资源不足 "
                f"cpu_usage={cpu_usage} gpu_usage={gpu_usage} "
                f"free_cpu={self.cpu_resource} free_gpu={self.gpu_resource}"
            )
            return False
        
    async def release(self, cpu_usage: int, gpu_usage: int) -> None:
        """ 释放资源 """
        before_cpu = self.cpu_resource
        before_gpu = self.gpu_resource
        self.cpu_resource += cpu_usage
        self.gpu_resource += gpu_usage
        logger.calculator.info(
            "[ResourceLocal] 释放本地资源 "
            f"cpu_usage={cpu_usage} gpu_usage={gpu_usage} "
            f"before_cpu={before_cpu} after_cpu={self.cpu_resource} "
            f"before_gpu={before_gpu} after_gpu={self.gpu_resource}"
        )
        await self.context.submit_event("calculator", Events.RELEASE_RESOURCE, {
            "cpu_resource": cpu_usage,
            "gpu_resource": gpu_usage,
        })
