#!/usr/bin/env python3
"""nvidia_api.py
Simple CLI wrapper for NVIDIA Inference API (OpenAI-compatible).

Usage:
  python3 nvidia_api.py health
  python3 nvidia_api.py chat --model <model> --prompt "..."
  python3 nvidia_api.py bench --model <model> --requests 3

The script looks for credentials in ~/.openclaw/credentials/nvidia.key or env NVIDIA_API_KEY.
"""
import os
import sys
import json
import time
import argparse
from typing import Optional

import requests

BASE = os.environ.get('NVIDIA_BASE_URL','https://integrate.api.nvidia.com/v1')
KEY = os.environ.get('NVIDIA_API_KEY')
if not KEY:
    key_path = os.path.expanduser('~/.openclaw/credentials/nvidia.key')
    if os.path.exists(key_path):
        with open(key_path,'r') as f:
            KEY = f.read().strip()

HEADERS = {
    'Authorization': f'Bearer {KEY}' if KEY else '',
    'Content-Type': 'application/json'
}


def health():
    url = f"{BASE}/models"
    r = requests.get(url, headers=HEADERS, timeout=20)
    try:
        r.raise_for_status()
    except Exception as e:
        print(json.dumps({'ok': False, 'error': str(e), 'status_code': r.status_code, 'body': r.text}))
        sys.exit(2)
    data = r.json()
    print(json.dumps({'ok': True, 'models_count': len(data.get('data', data) if isinstance(data, dict) else data)}))
    return data


def chat(model: str, prompt: str, temperature: float=0.2, max_tokens: int=512):
    url = f"{BASE}/chat/completions"
    payload = {
        'model': model,
        'messages': [{'role':'user','content': prompt}],
        'temperature': temperature,
        'max_tokens': max_tokens,
    }
    r = requests.post(url, headers=HEADERS, json=payload, timeout=60)
    try:
        r.raise_for_status()
    except Exception as e:
        print(json.dumps({'ok': False, 'error': str(e), 'status_code': r.status_code, 'body': r.text}))
        sys.exit(2)
    data = r.json()
    # Print concise output
    choice = None
    try:
        choice = data['choices'][0]['message']
    except Exception:
        choice = data
    print(json.dumps({'ok': True, 'response': choice}))
    return data


def bench(model: str, prompt: str = 'Hello', requests_n: int =3):
    results = []
    for i in range(requests_n):
        t0 = time.time()
        try:
            data = chat(model, prompt)
            ok = True
        except SystemExit:
            ok = False
            data = None
        dt = time.time() - t0
        results.append({'index': i, 'ok': ok, 'latency_s': dt})
        time.sleep(0.5)
    summary = {
        'model': model,
        'requests': requests_n,
        'avg_latency_s': sum(r['latency_s'] for r in results)/len(results) if results else None,
        'results': results
    }
    print(json.dumps({'ok': True, 'summary': summary}))
    return summary


def main():
    p = argparse.ArgumentParser()
    sub = p.add_subparsers(dest='cmd')
    sub.add_parser('health')
    pc = sub.add_parser('chat')
    pc.add_argument('--model', required=True)
    pc.add_argument('--prompt', required=True)
    pc.add_argument('--temperature', type=float, default=0.2)
    pc.add_argument('--max_tokens', type=int, default=512)
    pb = sub.add_parser('bench')
    pb.add_argument('--model', required=True)
    pb.add_argument('--prompt', default='Hello')
    pb.add_argument('--requests', type=int, default=3)

    args = p.parse_args()
    if args.cmd == 'health':
        health()
    elif args.cmd == 'chat':
        chat(args.model, args.prompt, temperature=args.temperature, max_tokens=args.max_tokens)
    elif args.cmd == 'bench':
        bench(args.model, args.prompt, requests_n=args.requests)
    else:
        p.print_help()

if __name__ == '__main__':
    main()