import time
import gpustat
import requests
import argparse
import os
import re

def get_mechineInfo():
    return [{
        'idx': i, 'temperature': card['temperature.gpu'], 
        'percent':card['utilization.gpu'], 'occupy':card['memory.used'], 
        'users': [
            {'name': p['username'], 'occupy': p['gpu_memory_usage']}
            for p in card['processes'] if p['gpu_memory_usage'] > 100
        ]}
        for i, card in enumerate(gpustat.new_query().jsonify()['gpus'])
    ]

# get args name ip port
parser = argparse.ArgumentParser()
parser.add_argument('--name', type=str, default='default')
parser.add_argument('--ip', type=str, default=None)
parser.add_argument('--port', type=int, default=5000)

args = parser.parse_args()

def get_server_host_from_index():
    current_path = os.path.dirname(os.path.abspath(__file__))
    with open(f'{current_path}/index.html', 'r') as f:
        content = f.read()
    p = re.compile(r'http://(.*?):\d+/gpu')
    match = p.search(content)
    return match.group(1)

server_host = get_server_host_from_index() if args.ip is None else args.ip
url = f"http://{server_host}:{args.port}/gpu/{args.name}"

while True:
    message = {
        'name': args.name,
        'gpuInfos': get_mechineInfo()
    }
    try:
        requests.post(url, json=message)
    except Exception as e:
        print(e)
    time.sleep(1)