#!/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 subprocess
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 key_cmds():
    """Готовые команды выдачи ключа (`key_cmd` в config.yaml) — по профилям.

    В сборке 1.7.x ключа в .env нет вообще: клиент получает короткоживущий
    токен `as-…` от gw-key.ps1, и шлюз принимает его на /key/info. Без этого
    чип всегда показывал «ключ шлюза не найден».
    """
    out = []
    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_cmd\s*:\s*(.+)$", text, re.M):
            val = m.group(1).strip()
            if val in ("", "null", "~") or val.startswith("#"):
                continue
            if len(val) > 1 and val[0] == val[-1] and val[0] in "'\"":
                val = val[1:-1].replace("''", "'")
            if val not in out:
                out.append(val)
    return out


def token_from_cmd(cmd):
    """Запуск key_cmd: единственная строка JSON, ключ — access_token."""
    kwargs = {}
    if os.name == "nt":
        kwargs["creationflags"] = getattr(subprocess, "CREATE_NO_WINDOW", 0)
    proc = subprocess.run(
        cmd, shell=True, capture_output=True, text=True, timeout=20, **kwargs
    )
    text = (proc.stdout or "").strip()
    if proc.returncode != 0 or not text:
        err = (proc.stderr or "").strip().splitlines()
        hint = err[-1][:120] if err else "код %s" % proc.returncode
        raise RuntimeError("key_cmd: %s" % hint)
    data = json.loads(text.splitlines()[-1])
    token = data.get("access_token") or data.get("key")
    if not token:
        raise RuntimeError("key_cmd: в ответе нет access_token")
    return token


def key_candidates():
    """(метка, ключ) — сначала .env, затем key_cmd профилей."""
    names = config_key_names()
    for name in KEY_NAMES:
        if name not in names:
            names.append(name)

    seen = set()
    for path in env_candidates():
        name, key = read_key(path, names)
        if key and key not in seen:
            seen.add(key)
            yield "%s из .env" % name, key

    for cmd in key_cmds():
        try:
            token = token_from_cmd(cmd)
        except Exception as e:
            raise RuntimeError(str(e)[:160])
        if token and token not in seen:
            seen.add(token)
            yield "токен gw-key", token


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 mark(state):
    """Последний запуск — короткой строкой в logs/ailab-limit-last.json.

    Проверка для поддержки: молчит чип = нет запусков, а не «нет расхода».
    """
    try:
        state = dict(state, ts=time.strftime("%Y-%m-%d %H:%M:%S"))
        path = os.path.join(root_dir(), "logs", "ailab-limit-last.json")
        with open(path, "w", encoding="utf-8") as f:
            f.write(json.dumps(state, ensure_ascii=False))
    except OSError:
        pass


def main():
    last_err = "ключ шлюза не найден: нет ни sk- в .env, ни key_cmd"

    try:
        candidates = list(key_candidates())
    except RuntimeError as e:
        mark({"ok": False, "err": str(e)})
        print(json.dumps({"ok": False, "err": str(e)}, ensure_ascii=False))
        return

    for label, key in candidates:
        try:
            info = fetch(key)
        except urllib.error.HTTPError as e:
            last_err = "шлюз вернул %s (%s)" % (e.code, label)
            continue
        except Exception as e:
            mark({"ok": False, "err": str(e)[:120]})
            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

        payload = {
            "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 "",
            "total_spend": round(float(info.get("total_spend") or 0), 2),
            "reset_at": info.get("budget_reset_at") or "",
            "source": label,
        }
        mark(payload)
        print(json.dumps(payload, ensure_ascii=False))
        return

    mark({"ok": False, "err": last_err})
    print(json.dumps({"ok": False, "err": last_err}, ensure_ascii=False))


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