#!/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 в рублях.

Имя переменной НЕ угадывается: скрипт читает `key_env` из config.yaml (активного
профиля, затем корневого) — у разных установок ключ лежит под разными именами
(AILAB, POLZA и т.п.), и жёсткий список молча давал «ключ не найден».

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

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

BASE = "https://gw.ailabscrume.ru"

# Запасной список: если config.yaml не найден или в нём нет key_env.
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",
)

# shell.exec рвёт вызов на 30 с, поэтому 3 попытки должны укладываться в этот
# бюджет с запасом: 3 x 7 с + паузы = 24 с.
TIMEOUT = 7.0
ATTEMPTS = 3
PAUSE = 1.5


def root_dir():
    return os.path.expandvars(r"%LOCALAPPDATA%\hermes")


def active_profile():
    path = os.path.expandvars(r"%APPDATA%\hermes\active-profile.json")
    if not os.path.exists(path):
        return None
    try:
        with open(path, encoding="utf-8") as f:
            data = json.load(f)
        return data.get("profile") or data.get("active")
    except Exception:
        return None


def profile_dirs():
    """Каталоги профилей: активный первым, затем остальные по алфавиту."""
    root = root_dir()
    out = []

    name = active_profile()
    if name:
        out.append(os.path.join(root, "profiles", name))

    out.append(root)

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

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


def config_key_names():
    """Имена переменных из key_env в config.yaml — любые, а не жёсткий список."""
    names = []
    for d in profile_dirs():
        cfg = os.path.join(d, "config.yaml")
        try:
            with open(cfg, encoding="utf-8", errors="ignore") as f:
                text = f.read()
        except OSError:
            continue
        for m in re.finditer(r"^\s*key_env\s*:\s*[\"\']?([A-Za-z0-9_]+)", text, re.M):
            name = m.group(1)
            if name not in names:
                names.append(name)
    return names


def env_candidates():
    """Пути к .env в том же порядке, что и профили."""
    return [os.path.join(d, ".env") for d in profile_dirs()]


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

    wanted = set(names)
    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 wanted:
            value = value.strip().strip('"').strip("'")
            if value.startswith("sk-"):
                found[name] = value

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


def fetch(key):
    """GET /key/info с повторами только на сетевых сбоях.

    401/403 — реальный отказ ключа, повторять бессмысленно: отдаём сразу.
    """
    last = None
    for attempt in range(ATTEMPTS):
        try:
            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)
        except urllib.error.HTTPError as e:
            if e.code in (401, 403):
                raise
            last = e
        except Exception as e:
            last = e
        if attempt + 1 < ATTEMPTS:
            time.sleep(PAUSE)
    raise last


def main():
    names = config_key_names()
    for name in KEY_NAMES:
        if name not in names:
            names.append(name)

    for path in env_candidates():
        name, key = read_key(path, names)
        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())