import json

import requests

from llm_server import opts


def get_power_states():
    gpu_num = 0
    output = {}
    while True:
        url = f"{opts.netdata_root}/api/v1/data?chart=nvidia_smi.gpu{gpu_num}_power_state"
        try:
            response = requests.get(url, timeout=3)
            if response.status_code != 200:
                break
            data = json.loads(response.text)
            power_state_data = data['data'][0]
            power_state = None
            for i in range(1, len(power_state_data)):
                if power_state_data[i] == 1:
                    power_state = data['labels'][i]
                    break
            output[f'gpu{gpu_num}'] = int(power_state.lower().strip('p'))
        except Exception as e:
            print('Failed to fetch Netdata metrics:', e)
            return output
        gpu_num += 1
    return output