#!/usr/bin/env python3.10

import subprocess

def get_accounts():
    command = ['sacctmgr', '--noheader', 'list', 'account', 'format=account']
    output = subprocess.run(command,capture_output=True,encoding='utf-8')
    accounts = output.stdout.split() 
    return accounts

def get_hours(account):
    command = ['sacct', '-X', '-a', '-A', str(account), 
            '--starttime=2025-01-01', '--parsable2', '--noheader', 
            '--format=elapsedraw,ncpus']
    cpu_hours_command = ['sreport', 'cluster', 'userutilizationbyaccount', 
            '--noheader', '--parsable2', '--tres=cpu', 
            'accounts='+str(account), 'format=used', 
            'start=2025-05-01T00:00:00']
    gpu_hours_command = ['sreport', 'cluster', 'userutilizationbyaccount', 
            '--noheader', '--parsable2', '--tres=gres/gpu', 
            'accounts='+str(account), 'format=used', 
            'start=2025-05-01T00:00:00']

    cpu_hours = subprocess.run(
            cpu_hours_command,
            capture_output=True,
            encoding='utf-8').stdout.split()
    
    cpu_hours = [int(i) for i in cpu_hours]
    cpu_hours = sum(cpu_hours)

    gpu_hours = subprocess.run(
            gpu_hours_command,
            capture_output=True,
            encoding='utf-8').stdout.split()

    gpu_hours = [int(i) for i in gpu_hours]
    gpu_hours = sum(gpu_hours)
    
    return cpu_hours, gpu_hours 

def update_qos(account):
    match account:
        case a if a.startswith('pad'):
            default_cpu, default_gpu = 240000000,15000000 
        case a if a.startswith('pci'):
            default_cpu, default_gpu = 60000000,3600000
        case a if a.startswith('pisca'):
            default_cpu, default_gpu = 6000000,60000 

    spent_cpu, spent_gpu = get_hours(account)

    new_cpu = max(0,default_cpu-spent_cpu)
    new_gpu = max(0,default_gpu-spent_gpu)
    
    update_command = ['sacctmgr', 'modify', 'qos', 'qos_'+str(account), 'set', 
            'GrpTRESMins=cpu='+str(new_cpu)+',gres/gpu='+str(new_gpu)]

    print(update_command)
    subprocess.run(update_command)

def update_defqos(account):
    '''
    Sets Default QOS to IPAC project in case of misconfiguration.
    '''
    command =  ['sacctmgr', 'modify', 'account', 'set', 
            'defaultqos=qos_'+str(account), 'where', 'account='+str(account)]
    subprocess.run(command)

def action(accounts,func=get_hours):
    for account in accounts:
        if account.startswith(('pad','pci','pisca')):
            func(account)

actions = {
        '1': get_hours,
        '2': update_defqos,
        '3': update_qos,
}

if __name__ == '__main__':
    print("Elegir que acción ejecutar en slurm:\n","1) get_hours\n",
            "2) update_defqos\n","3) update_qos\n")
    choice = input("Elegir:")
    func = actions.get(choice)

    accounts = get_accounts()
    if func:
        action(accounts,func)
    else:
        print("Operación invalida")