from src.manager.utils.singleton import singleton
import torch
import psutil

@singleton
class BudgetManager():
    total_resource_budget = 100
    current_resource_usage = 0
    total_expense_budget = 1000
    current_expense = 0
    is_budget_initialized = False
    
    def __init__(self):
        if not self.is_budget_initialized:
            self.total_resource_budget = self.calculate_total_budget()
            self.is_budget_initialized = True
        
    def calculate_total_budget(self)-> int:
        total_mem = 0
        if torch.cuda.is_available():
            gpu_index = torch.cuda.current_device()
            gpu_name = torch.cuda.get_device_name(gpu_index)
            total_vram = torch.cuda.get_device_properties(gpu_index).total_memory
            total_mem = total_vram /1024 ** 3
            print(f"GPU detected: {gpu_name}")
            print(f"Total VRAM: {total_mem:.2f} GB")
        else:
            mem = psutil.virtual_memory()
            total_mem = mem.total/ 1024 ** 3
            print("No GPU detected. Using CPU.")
            print(f"Total RAM: {total_mem:.2f} GB")
        return round((total_mem / 16) * 100)

    def get_total_resource_budget(self):
        return self.total_resource_budget
    
    def get_current_resource_usage(self):
        return self.current_resource_usage
    
    def get_current_remaining_resource_budget(self):
        return self.total_resource_budget - self.current_resource_usage
    
    def can_spend_resource(self, cost):
        return True if self.current_resource_usage + cost <= self.total_resource_budget else False
        
    def add_to_resource_budget(self, cost):
        if not self.can_spend_resource(cost):
            raise Exception("No resource budget remaining")
        self.current_resource_usage += cost
    
    def remove_from_resource_expense(self, cost):
        if self.current_resource_usage - cost < 0:
            raise Exception("Not enough resource budget to remove")
        self.current_resource_usage -= cost
    
    def get_total_expense_budget(self):
        return self.total_expense_budget
    
    def get_current_expense(self):
        return self.current_expense
    
    def get_current_remaining_expense_budget(self):
        return self.total_expense_budget - self.current_expense
    
    def can_spend_expense(self, cost):
        return True if self.current_expense + cost <= self.total_expense_budget else False
    
    def add_to_expense_budget(self, cost):
        if not self.can_spend_expense(cost):
            raise Exception("No expense budget remaining")
        self.current_expense += cost