#!/usr/bin/env python3
"""collect_usage.py
Collects usage & health metrics from configured providers and logs into SQLite.
Runs safely without sudo; writes DB to /home/alex/clawd/data/usage.db
"""
import os
import sqlite3
import time
import json
import subprocess
from datetime import datetime

DB_DIR = os.path.expanduser('~/clawd/data')
DB_PATH = os.path.join(DB_DIR, 'usage.db')
LOG_PATH = '/var/log/openclaw/usage.log'

os.makedirs(DB_DIR, exist_ok=True)

def init_db(conn):
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS metrics (
        ts INTEGER,
        provider TEXT,
        metric TEXT,
        value REAL,
        raw TEXT
    )''')
    conn.commit()

def insert(conn, provider, metric, value, raw=None):
    c = conn.cursor()
    c.execute('INSERT INTO metrics (ts,provider,metric,value,raw) VALUES (?,?,?,?,?)',
              (int(time.time()), provider, metric, value or 0.0, json.dumps(raw) if raw is not None else None))
    conn.commit()

def check_ollama():
    try:
        out = subprocess.check_output(['curl','-sS','http://192.168.1.144:11434/api/tags'], timeout=10)
        j = json.loads(out)
        models = len(j.get('models', []))
        return True, {'models': models}
    except Exception as e:
        return False, {'error': str(e)}

def check_nvidia():
    # call nvidia_api.py health if available
    script = os.path.expanduser('~/clawd/skills/nvidia-agent/scripts/nvidia_api.py')
    if os.path.exists(script):
        try:
            out = subprocess.check_output([script,'health'], timeout=20)
            j = json.loads(out.decode())
            return True, j
        except Exception as e:
            return False, {'error': str(e)}
    return None, None

def check_logs_for_rate_limits(window_minutes=30):
    res = []
    cutoff = time.time() - (window_minutes*60)
    if os.path.exists(LOG_PATH):
        try:
            with open(LOG_PATH,'r') as f:
                for line in f:
                    # Expect ISO timestamp at start
                    if not line.strip():
                        continue
                    try:
                        ts_part = line.split()[0]
                        try:
                            ts = datetime.fromisoformat(ts_part).timestamp()
                        except Exception:
                            ts = None
                    except Exception:
                        ts = None
                    if ts and ts < cutoff:
                        continue
                    if 'rate_limit' in line or '429' in line or 'RESOURCE_EXHAUSTED' in line:
                        res.append(line.strip())
        except Exception:
            pass
    return res


def main():
    conn = sqlite3.connect(DB_PATH)
    init_db(conn)

    # Ollama
    ok, data = check_ollama()
    if ok:
        insert(conn, 'ollama', 'models_count', data.get('models',0), data)
    else:
        insert(conn, 'ollama', 'up', 0, data)

    # NVIDIA
    ok, data = check_nvidia()
    if ok is True:
        insert(conn, 'nvidia', 'models_count', data.get('models_count',0), data)
    elif ok is False:
        insert(conn, 'nvidia', 'up', 0, data)

    # Logs rate limits
    rl = check_logs_for_rate_limits(30)
    insert(conn, 'logs', 'rate_limit_hits_30m', len(rl), {'lines': rl[:10]})

    # Optionally OpenAI admin cost (if key present)
    admin_key_path = os.path.expanduser('~/.openclaw/credentials/openai-admin.key')
    if os.path.exists(admin_key_path):
        try:
            key = open(admin_key_path).read().strip()
            import requests
            r = requests.get('https://api.openai.com/v1/usage', headers={'Authorization': f'Bearer {key}'}, timeout=10)
            if r.status_code == 200:
                js = r.json()
                # placeholder parsing
                insert(conn, 'openai', 'admin_raw_ok', 1, js)
            else:
                insert(conn, 'openai', 'admin_ok', 0, {'status': r.status_code, 'body': r.text})
        except Exception as e:
            insert(conn, 'openai', 'admin_ok', 0, {'error': str(e)})

    conn.close()

if __name__ == '__main__':
    main()