# GPU Roster Loader - reads gpu_roster.yaml via PyYAML
# Hot-reloadable via /admin/roster/reload endpoint

import os, json, threading, time

try:
    import yaml
except ImportError:
    yaml = None

ROSTER_PATH = os.environ.get('ROSTER_PATH', '/app/gpu_roster.yaml')
_last_mtime = 0
_lock = threading.Lock()

GPU_URLS = {}
GPU_SIDECARS = {}
GPU_LABELS = {}
GPU_MAX_CONCURRENT = {}
GPU_CONTEXT = {}
TIER_MODELS = {}
HOSTS = {}

def load_roster(path=None):
    global GPU_URLS, GPU_SIDECARS, GPU_LABELS, GPU_MAX_CONCURRENT
    global GPU_CONTEXT, TIER_MODELS, HOSTS, _last_mtime
    if yaml is None:
        return False, 'PyYAML not installed - run: pip3 install pyyaml'
    path = path or ROSTER_PATH
    try:
        with open(path) as f:
            data = yaml.safe_load(f)
        with _lock:
            GPU_URLS.clear()
            GPU_SIDECARS.clear()
            GPU_LABELS.clear()
            GPU_MAX_CONCURRENT.clear()
            GPU_CONTEXT.clear()
            TIER_MODELS.clear()
            models = data.get('models', {})
            for name, cfg in models.items():
                GPU_URLS[name] = cfg.get('gpu_url', '')
                GPU_SIDECARS[name] = cfg.get('sidecar_url', '')
                GPU_LABELS[name] = cfg.get('label', name)
                GPU_MAX_CONCURRENT[name] = cfg.get('max_concurrent', 1)
                GPU_CONTEXT[name] = cfg.get('context', 65536)
                tiers = cfg.get('tiers', ['enterprise'])
                for t in tiers:
                    if t not in TIER_MODELS:
                        TIER_MODELS[t] = []
                    if name not in TIER_MODELS[t]:
                        TIER_MODELS[t].append(name)
            HOSTS.clear()
            HOSTS.update(data.get('hosts', {}))
            _last_mtime = os.path.getmtime(path)
        return True, 'Loaded {} models, {} hosts'.format(len(GPU_URLS), len(HOSTS))
    except Exception as e:
        return False, str(e)

def check_reload():
    global _last_mtime
    try:
        mtime = os.path.getmtime(ROSTER_PATH)
        if mtime > _last_mtime:
            success, msg = load_roster()
            if success:
                print('[ROSTER] Auto-reloaded:', msg)
    except:
        pass

def reload_thread(interval=30):
    while True:
        time.sleep(interval)
        check_reload()

threading.Thread(target=reload_thread, daemon=True).start()