#!/usr/bin/env python3
"""Расход ключа шлюза AILab SCRUME (gw.ailabscrume.ru) — компактный JSON в stdout.

Вызывается плагином ailab-limit через shell.exec. Ключ берётся из .env профилей
и в рендерер не попадает. Один запрос: GET /key/info с пользовательским ключом
(master не нужен) — шлюз отдаёт spend, max_budget и budget_duration в рублях.

Формат ответа:
  {"ok":true,"spend":5.21,"budget":100.0,"duration":"1d","pct":5.2,"alias":"testoff"}
  {"ok":false,"err":"..."}
"""

import json
import os
import sys
import urllib.error
import urllib.request

BASE = "https://gw.ailabscrume.ru"
KEY_NAMES = (
    "HERMES_CUSTOM_AILAB_API_KEY",
    "CUSTOM_PROVIDER_AILAB_KEY",
    "HERMES_CUSTOM_POLZA_AI_API_KEY",
    "CUSTOM_PROVIDER_POLZA_AI_KEY",
    "AILAB_API_KEY",
)


def candidates():
    """Пути к .env: корневой и все профили, активный — первым."""
    root = os.path.expandvars(r"%LOCALAPPDATA%\hermes")
    out = []

    active = os.path.expandvars(r"%APPDATA%\hermes\active-profile.json")
    if os.path.exists(active):
        try:
            with open(active, encoding="utf-8") as f:
                name = json.load(f).get("profile") or json.load(f).get("active")
            if name:
                out.append(os.path.join(root, "profiles", name, ".env"))
        except Exception:
            pass

    out.append(os.path.join(root, ".env"))
    prof_dir = os.path.join(root, "profiles")
    if os.path.isdir(prof_dir):
        for name in sorted(os.listdir(prof_dir)):
            out.append(os.path.join(prof_dir, name, ".env"))

    seen, uniq = set(), []
    for p in out:
        if p not in seen:
            seen.add(p)
            uniq.append(p)
    return uniq


def read_key(path):
    """Первая непустая переменная из KEY_NAMES в файле .env."""
    try:
        with open(path, encoding="utf-8", errors="ignore") as f:
            lines = f.readlines()
    except OSError:
        return None, None

    found = {}
    for line in lines:
        s = line.strip()
        if not s or s.startswith("#") or "=" not in s:
            continue
        name, _, value = s.partition("=")
        name = name.strip()
        if name in KEY_NAMES:
            value = value.strip().strip('"').strip("'")
            if value.startswith("sk-"):
                found[name] = value

    for name in KEY_NAMES:
        if name in found:
            return name, found[name]
    return None, None


def fetch(key, timeout=10.0):
    req = urllib.request.Request(BASE + "/key/info", headers={"Authorization": "Bearer " + key})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        data = json.loads(r.read())
    return data.get("info", data)


def main():
    for path in candidates():
        name, key = read_key(path)
        if not key:
            continue
        try:
            info = fetch(key)
        except urllib.error.HTTPError as e:
            print(json.dumps({"ok": False, "err": "шлюз вернул %s" % e.code}, ensure_ascii=False))
            return
        except Exception as e:
            print(json.dumps({"ok": False, "err": str(e)[:120]}, ensure_ascii=False))
            return

        spend = float(info.get("spend") or 0)
        budget = info.get("max_budget")
        budget = float(budget) if budget is not None else None
        pct = round(spend / budget * 100, 1) if budget else None

        print(json.dumps({
            "ok": True,
            "spend": round(spend, 4),
            "budget": round(budget, 2) if budget else None,
            "duration": info.get("budget_duration") or "",
            "pct": pct,
            "alias": info.get("key_alias") or "",
        }, ensure_ascii=False))
        return

    print(json.dumps({"ok": False, "err": "ключ шлюза не найден в .env"}, ensure_ascii=False))


if __name__ == "__main__":
    sys.exit(main())