#!/usr/bin/env python3
"""Backend per il plasmoide KDE Plasma 6: Antigravity Usage Dashboard.

Interroga il gateway Google CloudCode PA (https://daily-cloudcode-pa.googleapis.com)
leggendo il token OAuth da ~/.gemini/antigravity-cli/antigravity-oauth-token
ed eseguendo il refresh automatico se scaduto o non valido.
Restituisce lo stato delle quote per le 3 famiglie di modelli:
- Gemini (Gemini 3.x Flash / Pro)
- Claude (Claude Sonnet 4.6 / Opus 4.6)
- GPT-OSS (GPT-OSS 120B)
"""

import json
import os
import sys
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone

TOKEN_FILE = os.path.expanduser("~/.gemini/antigravity-cli/antigravity-oauth-token")
OAUTH_URL = "https://oauth2.googleapis.com/token"
BASE_URL = "https://daily-cloudcode-pa.googleapis.com"
CLIENT_ID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
CLIENT_SECRET = "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"
UA = "antigravity/cli/1.1.17"
API_CLIENT = "google-cloud-sdk vscode_cloudshelleditor/0.1"
METADATA = json.dumps({"ideType": "ANTIGRAVITY", "platform": "LINUX", "pluginType": "GEMINI"})


def get_token_data():
    if not os.path.exists(TOKEN_FILE):
        return None
    try:
        with open(TOKEN_FILE, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None


def save_token_data(data):
    try:
        tmp = TOKEN_FILE + ".tmp"
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)
        os.replace(tmp, TOKEN_FILE)
    except Exception:
        pass


def refresh_token(refresh_tok):
    req_body = json.dumps({
        "client_id": CLIENT_ID,
        "client_secret": CLIENT_SECRET,
        "refresh_token": refresh_tok,
        "grant_type": "refresh_token"
    }).encode("utf-8")
    req = urllib.request.Request(OAUTH_URL, data=req_body, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=15) as resp:
        res = json.loads(resp.read().decode("utf-8"))
        return res.get("access_token"), int(res.get("expires_in", 3600))


def call_api(path, token, payload):
    url = f"{BASE_URL}{path}"
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
        "User-Agent": UA,
        "X-Goog-Api-Client": API_CLIENT,
        "Client-Metadata": METADATA
    }
    body = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(url, data=body, headers=headers)
    with urllib.request.urlopen(req, timeout=15) as resp:
        return json.loads(resp.read().decode("utf-8"))


def parse_iso_epoch(ts):
    if not ts:
        return 0
    try:
        dt = datetime.fromisoformat(ts.replace("Z", "+00:00"))
        return int(dt.timestamp())
    except Exception:
        return 0


def main():
    data = get_token_data()
    if not data or "token" not in data:
        print(json.dumps({
            "ok": False,
            "error": "no_token",
            "detail": f"Nessun token OAuth in {TOKEN_FILE}"
        }))
        sys.exit(0)

    tok = data["token"]
    access_tok = tok.get("access_token")
    refresh_tok = tok.get("refresh_token")

    project_id = "healthy-topic-dwr0v"
    try:
        quota_data = call_api("/v1internal:retrieveUserQuota", access_tok, {"project": project_id})
    except urllib.error.HTTPError as e:
        if e.code in (401, 403) and refresh_tok:
            try:
                new_acc, exp_in = refresh_token(refresh_tok)
                tok["access_token"] = new_acc
                save_token_data(data)
                access_tok = new_acc
                quota_data = call_api("/v1internal:retrieveUserQuota", access_tok, {"project": project_id})
            except Exception as e2:
                print(json.dumps({
                    "ok": False,
                    "error": "auth_error",
                    "detail": f"Refresh fallito: {e2}"
                }))
                sys.exit(0)
        else:
            print(json.dumps({
                "ok": False,
                "error": f"http_{e.code}",
                "detail": f"Errore server Google HTTP {e.code}"
            }))
            sys.exit(0)
    except Exception as e:
        print(json.dumps({
            "ok": False,
            "error": "network_error",
            "detail": str(e)
        }))
        sys.exit(0)

    buckets = quota_data.get("buckets", [])
    bucket_map = {b.get("modelId"): b for b in buckets}

    # 1. Gemini Family
    gem_b = (bucket_map.get("gemini-3.1-pro-high") or
             bucket_map.get("gemini-3.7-flash-high") or
             bucket_map.get("gemini-3.5-flash-low") or {})
    gem_rem = float(gem_b.get("remainingFraction", 1.0))
    gem_reset_str = gem_b.get("resetTime", "")
    gemini_info = {
        "id": "gemini",
        "name": "Gemini 3.x",
        "models": "Pro / Flash",
        "remaining": round(gem_rem, 4),
        "used_percent": round((1.0 - gem_rem) * 100, 1),
        "reset_time": gem_reset_str,
        "reset_at": parse_iso_epoch(gem_reset_str)
    }

    # 2. Claude Family
    cld_b = (bucket_map.get("claude-sonnet-4-6") or
             bucket_map.get("claude-opus-4-6-thinking") or {})
    cld_rem = float(cld_b.get("remainingFraction", 1.0))
    cld_reset_str = cld_b.get("resetTime", "")
    claude_info = {
        "id": "claude",
        "name": "Claude 4.6",
        "models": "Sonnet / Opus",
        "remaining": round(cld_rem, 4),
        "used_percent": round((1.0 - cld_rem) * 100, 1),
        "reset_time": cld_reset_str,
        "reset_at": parse_iso_epoch(cld_reset_str)
    }

    # 3. GPT-OSS Family
    gpt_b = bucket_map.get("gpt-oss-120b-medium") or {}
    gpt_rem = float(gpt_b.get("remainingFraction", 1.0))
    gpt_reset_str = gpt_b.get("resetTime", "")
    gpt_info = {
        "id": "gpt-oss",
        "name": "GPT-OSS",
        "models": "120B Medium",
        "remaining": round(gpt_rem, 4),
        "used_percent": round((1.0 - gpt_rem) * 100, 1),
        "reset_time": gpt_reset_str,
        "reset_at": parse_iso_epoch(gpt_reset_str)
    }

    out = {
        "ok": True,
        "project": project_id,
        "groups": [gemini_info, claude_info, gpt_info]
    }
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()