#!/usr/bin/env python3
"""
NVIDIA APIs Helper — wrapper unificado para build.nvidia.com
Carrega NVIDIA_API_KEY de scripts/.env
"""
import os
import sys
import json
import argparse
import urllib.request
import urllib.error
from pathlib import Path

ENV_FILE = Path(__file__).parent / ".env"
if ENV_FILE.exists():
    for line in ENV_FILE.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if "=" in line:
            k, v = line.split("=", 1)
            v = v.strip().strip('"').strip("'")
            os.environ.setdefault(k, v)

API_KEY = os.environ.get("NVIDIA_API_KEY", "")
BASE = "https://integrate.api.nvidia.com/v1"


def call_chat(model: str, prompt: str, max_tokens: int = 200, json_mode: bool = False) -> dict:
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
    }
    if json_mode:
        payload["response_format"] = {"type": "json_object"}

    req = urllib.request.Request(
        f"{BASE}/chat/completions",
        data=json.dumps(payload).encode(),
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}",
        },
    )
    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            return json.loads(resp.read().decode())
    except urllib.error.HTTPError as e:
        raw = e.read().decode()
        try:
            return json.loads(raw)
        except Exception:
            return {"_error": e.code, "_raw": raw[:500]}


def cmd_chat(args):
    if not API_KEY:
        print("ERRO: NVIDIA_API_KEY nao definida em scripts/.env", file=sys.stderr)
        print("Obtenha gratis em: https://build.nvidia.com/", file=sys.stderr)
        sys.exit(1)
    r = call_chat(args.model, args.prompt, args.max_tokens, args.json)
    if "_error" in r:
        print(f"ERRO {r['_error']}: {r['_raw']}")
        sys.exit(1)
    print(r["choices"][0]["message"]["content"])


def cmd_models(args):
    """Lista modelos disponiveis (sem auth)."""
    req = urllib.request.Request(f"{BASE}/models")
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            data = json.loads(resp.read().decode())
        models = [m["id"] for m in data.get("data", [])]
        print(f"{len(models)} modelos:")
        for m in sorted(models):
            if any(k in m.lower() for k in ["nemotron", "parakeet", "magpie", "nv-embed", "cosmos", "deepseek", "kimi", "mistral", "llama"]):
                print(f"  - {m}")
    except urllib.error.HTTPError as e:
        print(f"HTTP {e.code}: {e.read().decode()[:300]}")


def cmd_status(args):
    """Verifica se a key funciona."""
    if not API_KEY:
        print("NVIDIA_API_KEY: NAO configurada")
        print("Acao: obter gratis em https://build.nvidia.com/")
        return
    print(f"NVIDIA_API_KEY: configurada (prefixo: {API_KEY[:12]}...)")
    r = call_chat("nvidia/nemotron-mini-4b-instruct", "Diga 'ok' em portugues", 20)
    if "_error" in r:
        print(f"ERRO {r['_error']}: {r['_raw'][:200]}")
    else:
        print(f"Status: OK")
        print(f"Teste: {r['choices'][0]['message']['content']}")


def main():
    parser = argparse.ArgumentParser(description="NVIDIA APIs Helper (build.nvidia.com)")
    sub = parser.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("chat", help="chat completion")
    p.add_argument("prompt")
    p.add_argument("--model", default="nvidia/nemotron-mini-4b-instruct")
    p.add_argument("--max-tokens", type=int, default=200)
    p.add_argument("--json", action="store_true", help="force JSON output")
    p.set_defaults(func=cmd_chat)

    sub.add_parser("models", help="lista modelos disponiveis").set_defaults(func=cmd_models)

    sub.add_parser("status", help="verifica se key funciona").set_defaults(func=cmd_status)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()