#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""阿里云百炼 Token Plan 个人版额度探针（仅 Python 3 标准库，零依赖）。

两种模式（BAILIAN_MODE）：
  cookie   主路：走百炼控制台 OneConsole 网关，鉴权是「控制台 Cookie」。
           这是目前唯一被第三方实锤过的个人版 5 小时 / 7 天滚动窗口数据源
           （端点与字段对齐 token-monitor 项目从真实账号抓取的契约）。
           Cookie 随会话过期（通常几天），过期时快照标 error_kind=unauthorized。
  apikey   验证路：用 sk-sp API Key 打 Token Plan 网关上的候选 usage 端点并
           逐条记录响应。该路径存在性尚未证实（无 key 时网关一律先回
           InvalidApiKey，无法从 401 区分路径有无）；若某天命中窗口字段，
           快照与 cookie 模式完全同构，前端零改动。

输出契约：原子写 JSON 快照到 BAILIAN_OUT（前端只读该文件，绝不落 Cookie/Key）。
失败策略：不丢旧数据——从状态文件回填上次成功值并标 stale=true。

配置：环境变量，或 BAILIAN_PROBE_ENV 指向的 KEY=VALUE 文件（默认
/etc/bailian-probe.env，权限请设 600）。
"""

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

UA = ("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
      "(KHTML, like Gecko) Chrome/126.0 Safari/537.36")

# 个人版控制台变体（对齐 token-monitor 抓取契约；国内站默认）
VARIANTS = {
    "cn-personal": {
        "gateway_origin": "https://bailian.console.aliyun.com",
        "quota_origin": "https://bailian-cs.console.aliyun.com",
        "region_id": "cn-beijing",
        "commodity": "sfm_tokenplansolo_public_cn",
        "console_site": "BAILIAN_ALIYUN",
        "action": "BroadScopeAspnGateway",
    },
    "intl-personal": {
        "gateway_origin": "https://modelstudio.console.alibabacloud.com",
        "quota_origin": "https://bailian-singapore-cs.alibabacloud.com",
        "region_id": "ap-southeast-1",
        "commodity": "sfm_tokenplansolo_public_intl",
        "console_site": "MODELSTUDIO_ALBABACLOUD",  # 官方契约的历史拼写，保留
        "action": "IntlBroadScopeAspnGateway",
    },
}
PLAN_LABELS = {"lite": "Lite", "standard": "Standard", "pro": "Pro", "max": "Max"}

API_USAGE = "zeldaHttp.apikeyMgr./tokenplan/personal/api/v2/usage"
API_SUBSCRIPTION = "zeldaHttp.apikeyMgr./tokenplan/personal/api/v2/subscription"
API_QUOTA_CONFIG = "zeldaHttp.apikeyMgr./tokenplan/personal/api/v2/quota-config"
# 只读列表；reset-card/use 是真正消耗卡的写操作，探针永远不得调用
API_RESET_CARDS = "zeldaHttp.apikeyMgr./tokenplan/personal/api/v2/reset-card/list"

# apikey 模式的候选端点（存在性未证实，逐个探测并记录，命中即出图）
APIKEY_CANDIDATES = [
    "{base}/compatible-mode/v1/usage",
    "{base}/compatible-mode/v1/tokenplan/usage",
    "{base}/compatible-mode/v1/subscription/stats",
    "https://dashscope.aliyuncs.com/api/v1/tokenplan/subscription/stats",
]


def now_iso():
    return datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds")


def load_env_file(path):
    """读 KEY=VALUE 配置（不覆盖已有环境变量）。"""
    if not path or not os.path.exists(path):
        return
    with open(path, "r", encoding="utf-8-sig") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, v = line.split("=", 1)
            os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))


class ProbeError(Exception):
    """kind: unauthorized（凭证失效，需人工重贴）/ unavailable（临时不可用）。"""

    def __init__(self, kind, msg, diag=None):
        super().__init__(msg)
        self.kind = kind
        self.diag = diag


# --------------------------------------------------------------------------
# 通用小工具：数值 / 时间 / 全树大小写不敏感取值
# --------------------------------------------------------------------------

def expand_embedded(value):
    """网关常把 JSON 再字符串套一层，递归展开（token-monitor 的实锤教训）。"""
    if isinstance(value, str):
        s = value.strip()
        if s.startswith("{") or s.startswith("["):
            try:
                return expand_embedded(json.loads(s))
            except Exception:
                return value
        return value
    if isinstance(value, list):
        return [expand_embedded(v) for v in value]
    if isinstance(value, dict):
        return {k: expand_embedded(v) for k, v in value.items()}
    return value


def find_key(value, name):
    """全树、大小写不敏感找第一个非空值（网关字段拼写随变体/版本漂移）。"""
    wanted = name.lower()
    if isinstance(value, dict):
        for k, v in value.items():
            if k.lower() == wanted and v not in (None, "", []):
                return v
        for v in value.values():
            got = find_key(v, name)
            if got is not None:
                return got
    elif isinstance(value, list):
        for item in value:
            got = find_key(item, name)
            if got is not None:
                return got
    return None


def to_number(v):
    try:
        n = float(v)
        return n if n == n else None
    except (TypeError, ValueError):
        return None


def to_int_or_none(v):
    n = to_number(v)
    return int(n) if n is not None else None


def to_iso(v):
    """epoch 秒/毫秒/无时区字符串（按东八区钉死）统一转 ISO，防 8 小时偏移。"""
    n = to_number(v)
    if n is not None and n > 0:
        ms = n if n >= 1e12 else n * 1000
        return datetime.fromtimestamp(ms / 1000, timezone.utc).isoformat()
    if isinstance(v, str):
        s = v.strip().replace(" ", "T")
        if re.match(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}", s) and "+" not in s and not s.endswith("Z"):
            s += "+08:00"  # 控制台裸时间戳按北京时间解
        try:
            return datetime.fromisoformat(s).astimezone(timezone.utc).isoformat()
        except ValueError:
            return None
    return None


def pct_from_ratio(v):
    """网关口径为 0~1 占比；>1 时视为已是百分数，避免再放大 100 倍。"""
    n = to_number(v)
    if n is None:
        return None
    if n > 1:
        return round(min(max(n, 0), 100), 1)
    return round(min(max(n, 0), 1) * 100, 1)


def http(url, method="GET", headers=None, data=None, timeout=20):
    req = urllib.request.Request(url, method=method, data=data, headers=headers or {})
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, resp.read().decode("utf-8", "replace")
    except urllib.error.HTTPError as e:
        return e.code, e.read().decode("utf-8", "replace")


def cookie_value(cookie_header, name):
    for part in (cookie_header or "").split(";"):
        k, _, v = part.strip().partition("=")
        if k == name:
            return v.strip()
    return ""


# --------------------------------------------------------------------------
# 控制台网关（cookie 模式）
# --------------------------------------------------------------------------

def console_headers(cookie, variant):
    h = {
        "Accept": "application/json, text/plain, */*",
        "Content-Type": "application/x-www-form-urlencoded",
        "Cookie": cookie,
        "Origin": variant["gateway_origin"],
        "Referer": variant["gateway_origin"] + "/",
        "User-Agent": UA,
        "X-Requested-With": "XMLHttpRequest",
    }
    csrf = cookie_value(cookie, "login_aliyunid_csrf") or cookie_value(cookie, "csrf")
    if csrf:
        h["x-xsrf-token"] = csrf
        h["x-csrf-token"] = csrf
    return h


def resolve_sec_token(cookie, variant):
    """best-effort 解析 sec_token：JSON 端点 → cookie 兜底；拿不到不算错。"""
    try:
        status, text = http(variant["gateway_origin"] + "/tool/user/info.json",
                            headers={"Accept": "application/json", "Cookie": cookie,
                                     "User-Agent": UA,
                                     "Referer": variant["gateway_origin"] + "/"},
                            timeout=10)
        if status == 200:
            got = find_key(expand_embedded(json.loads(text)), "secToken")
            if got:
                return str(got)
    except Exception:
        pass
    return cookie_value(cookie, "sec_token")


def classify_body(payload, text):
    """网关用 HTTP 200 装一切错误：登录失效与工作空间未授权必须分开归因。"""
    low = (text or "").lower()
    if "<html" in low and "login" in low:
        raise ProbeError("unauthorized", "控制台会话已失效（返回登录页），请重新粘贴 Cookie")
    err_code = find_key(payload, "errorCode") or ""
    failed = (payload.get("success") is False or payload.get("Success") is False
              or bool(err_code))
    if failed:
        msg = str(find_key(payload, "errorMsg") or find_key(payload, "message") or err_code)
        joined = f"{err_code} {msg}".lower()
        if "workspace" in joined:
            # 工作空间授权问题 ≠ Cookie 失效，归因错了会诱导用户白换 Cookie
            raise ProbeError("unavailable", f"工作空间未授权（Cookie 本身有效）：{err_code or msg}")
        if any(k in joined for k in ("notlogined", "needlogin", "tokenerror",
                                     "unauthorized", "forbidden", "过期")):
            raise ProbeError("unauthorized", f"鉴权失败：{err_code or msg}")
        raise ProbeError("unavailable", f"网关错误：{err_code or msg}")


def fetch_personal_api(cookie, variant, api, sec_token, data_params):
    cornerstone = {
        "feTraceId": str(uuid.uuid4()),
        "feURL": variant["gateway_origin"] + "/" ,
        "protocol": "V2", "console": "ONE_CONSOLE", "productCode": "p_efm",
        "switchUserType": 3,
        "domain": urllib.parse.urlsplit(variant["gateway_origin"]).netloc,
        "consoleSite": variant["console_site"],
        "userNickName": "", "userPrincipalName": "", "xsp_lang": "zh-CN",
    }
    anon = cookie_value(cookie, "cna")
    if anon:
        cornerstone["X-Anonymous-Id"] = anon
    # 刻意不带 switchAgent：会把请求钉死在抓取账号的工作空间上
    params = {"Api": api, "V": "1.0",
              "Data": dict(data_params, cornerstoneParam=cornerstone)}
    body = {
        "product": "sfm_bailian", "action": variant["action"],
        "region": variant["region_id"], "language": "zh-CN",
        "params": json.dumps(params, ensure_ascii=False),
    }
    if sec_token:
        body["sec_token"] = sec_token
    qs = urllib.parse.urlencode({"action": variant["action"], "product": "sfm_bailian",
                                 "api": api, "_v": "undefined"})
    url = variant["quota_origin"] + "/data/api.json?" + qs
    status, text = http(url, "POST", console_headers(cookie, variant),
                        urllib.parse.urlencode(body).encode())
    if status in (401, 403):
        raise ProbeError("unauthorized", f"HTTP {status}，Cookie 失效")
    if status != 200:
        raise ProbeError("unavailable", f"HTTP {status}")
    try:
        payload = expand_embedded(json.loads(text))
    except Exception:
        raise ProbeError("unavailable", "网关返回非 JSON")
    classify_body(payload, text)
    return payload


def build_windows(five_raw, week_raw, five_reset_raw, week_reset_raw, totals):
    """两条窗口的共用组装（cookie / apikey 模式同构输出）。"""
    windows = []
    five_pct = pct_from_ratio(five_raw)
    week_pct = pct_from_ratio(week_raw)
    if five_pct is not None:
        ft = to_number(totals.get("five_hour") or totals.get("fiveHour"))
        windows.append({
            "id": "5h", "label": "5 小时", "used_percent": five_pct,
            "total": ft,
            "remaining": round(ft * (100 - five_pct) / 100, 1) if ft else None,
            "reset_at": to_iso(five_reset_raw),
        })
    if week_pct is not None:
        wt = to_number(totals.get("weekly"))
        windows.append({
            "id": "weekly", "label": "7 天", "used_percent": week_pct,
            "total": wt,
            "remaining": round(wt * (100 - week_pct) / 100, 1) if wt else None,
            "reset_at": to_iso(week_reset_raw),
        })
    return windows


def probe_cookie_mode(variant_id):
    cookie = os.environ.get("BAILIAN_COOKIE", "")
    if not cookie:
        raise ProbeError("unauthorized", "未配置 BAILIAN_COOKIE")
    variant = VARIANTS[variant_id]
    sec = resolve_sec_token(cookie, variant)

    # usage 网关偶发 200 但不带窗口字段：官方行为，最多 3 次（对齐上游实现）
    usage = None
    for _ in range(3):
        usage = fetch_personal_api(cookie, variant, API_USAGE, sec, {})
        if find_key(usage, "per5HourPercentage") or find_key(usage, "per1WeekPercentage"):
            break
        time.sleep(0.4)
    five = find_key(usage, "per5HourPercentage")
    week = find_key(usage, "per1WeekPercentage")
    if five is None and week is None:
        raise ProbeError("unavailable", "网关未返回滚动窗口字段")

    # 订阅档位与配额上限：best-effort，丢只丢绝对值，不影响百分比
    spec_code, totals, subscription = "", {}, None
    try:
        sub = fetch_personal_api(cookie, variant, API_SUBSCRIPTION, sec,
                                 {"commodityCode": variant["commodity"]})
        spec_code = str(find_key(sub, "specCode") or find_key(sub, "spec_code") or "").lower()
        subscription = {
            "status": str(find_key(sub, "status") or "").upper() or None,
            "remaining_days": to_int_or_none(find_key(sub, "remainingDays")),
            "start_at": to_iso(find_key(sub, "startTime")),
            "end_at": to_iso(find_key(sub, "endTime")),
            "auto_renew": bool(find_key(sub, "autoRenewFlag")) if find_key(sub, "autoRenewFlag") is not None else None,
        }
    except ProbeError:
        pass
    try:
        qc = fetch_personal_api(cookie, variant, API_QUOTA_CONFIG, sec, {})
        if spec_code:
            node = find_key(qc, spec_code)
            if isinstance(node, dict):
                totals = node
    except ProbeError:
        pass

    windows = build_windows(five, week,
                            find_key(usage, "per5HourResetTime"),
                            find_key(usage, "per1WeekResetTime"), totals)
    reset_cards = fetch_reset_cards(cookie, variant, sec)
    diag = {"usage_keys": sorted(usage.keys()) if isinstance(usage, dict) else [],
            "quota_totals_found": bool(totals),
            "reset_cards_found": len(reset_cards)}
    return windows, PLAN_LABELS.get(spec_code, spec_code or "Personal"), diag, reset_cards, subscription


# 重置卡类型 → 中文用途（目前控制台只有 1W 一种，未知类型保留原码）
RESET_CARD_LABELS = {"RESET_1W": "7 天窗口重置卡"}


def extract_card_list(payload):
    """reset-card/list 的业务数据是嵌了三层 data 的数组，按结构特征（含 cardNo）定位，
    比写死路径更抗网关版本漂移。"""
    found = []

    def walk(o):
        if isinstance(o, list):
            if o and all(isinstance(x, dict) and "cardNo" in x for x in o):
                found.append(o)
            for x in o:
                walk(x)
        elif isinstance(o, dict):
            for x in o.values():
                walk(x)

    walk(payload)
    return found[0] if found else []


def fetch_reset_cards(cookie, variant, sec_token):
    """best-effort：卡列表是附属信息，接口异常绝不影响主用量快照。"""
    try:
        raw = fetch_personal_api(cookie, variant, API_RESET_CARDS, sec_token, {})
        cards = []
        for c in extract_card_list(raw):
            ctype = str(c.get("cardType") or "")
            cards.append({
                "type": ctype,
                "label": RESET_CARD_LABELS.get(ctype, ctype or "重置卡"),
                "card_no": str(c.get("cardNo") or ""),
                "effective_at": to_iso(c.get("effectiveAt")),
                "expires_at": to_iso(c.get("expiresAt")),
            })
        return cards
    except Exception:
        return []


# --------------------------------------------------------------------------
# API Key 候选端点（apikey 模式，验证性质）
# --------------------------------------------------------------------------

def probe_apikey_mode():
    key = os.environ.get("BAILIAN_API_KEY", "")
    base = os.environ.get("BAILIAN_BASE_URL",
                          "https://token-plan.cn-beijing.maas.aliyuncs.com")
    if not key:
        raise ProbeError("unauthorized", "未配置 BAILIAN_API_KEY")
    results, windows = [], []
    for tpl in APIKEY_CANDIDATES:
        url = tpl.format(base=base.rstrip("/"))
        try:
            status, text = http(url, headers={"Authorization": "Bearer " + key,
                                              "Accept": "application/json",
                                              "User-Agent": UA}, timeout=15)
        except Exception as e:
            results.append({"url": url, "error": str(e)[:200]})
            continue
        entry = {"url": url, "status": status, "excerpt": text[:300]}
        try:
            payload = expand_embedded(json.loads(text))
            f = find_key(payload, "per5HourPercentage")
            w = find_key(payload, "per1WeekPercentage")
            entry["window_like"] = f is not None or w is not None
            if entry["window_like"] and not windows:
                windows = build_windows(f, w,
                                        find_key(payload, "per5HourResetTime"),
                                        find_key(payload, "per1WeekResetTime"), {})
        except Exception:
            entry["window_like"] = False
        results.append(entry)
    if not windows:
        # 全部未命中：探测记录仍要进快照 diagnostics，这就是本模式的全部意义
        raise ProbeError("unavailable", "候选 usage 端点均未返回窗口字段（见 diagnostics）",
                         diag={"apikey_probe": results})
    plan = os.environ.get("BAILIAN_PLAN_NAME", "Personal")
    return windows, plan, {"apikey_probe": results}


# --------------------------------------------------------------------------
# 快照与状态
# --------------------------------------------------------------------------

def read_json(path):
    try:
        # utf-8-sig：容忍 Windows 记事本/PowerShell 写出的带 BOM 文件
        with open(path, "r", encoding="utf-8-sig") as f:
            return json.load(f)
    except Exception:
        return None


def atomic_write(path, obj):
    parent = os.path.dirname(os.path.abspath(path))
    os.makedirs(parent, exist_ok=True)
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(obj, f, ensure_ascii=False, indent=1)
    os.replace(tmp, path)


def zero_without_rollover(new_windows, last_windows):
    """窗口没滚出（reset_at 未变早）却从高位掉回 0%：语义可疑，只标记不改数。"""
    flags = []
    for w in new_windows:
        prev = next((x for x in last_windows or [] if x.get("id") == w.get("id")), None)
        if (w.get("used_percent") == 0 and prev and (prev.get("used_percent") or 0) >= 5
                and w.get("reset_at") and w.get("reset_at") == prev.get("reset_at")):
            # 字符串直接比较会因时区偏移不同而失真，按 datetime 比较
            try:
                if datetime.fromisoformat(w["reset_at"]) > datetime.now(timezone.utc):
                    flags.append(w["id"])
            except ValueError:
                pass
    return flags


def main():
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    load_env_file(os.environ.get("BAILIAN_PROBE_ENV", "/etc/bailian-probe.env"))
    out = os.environ.get("BAILIAN_OUT",
                         os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                      "bailian.json"))
    state_path = os.environ.get("BAILIAN_STATE", out + ".state.json")
    log_path = os.environ.get("BAILIAN_LOG", out + ".probe.jsonl")
    mode = os.environ.get("BAILIAN_MODE", "cookie")
    variant_id = os.environ.get("BAILIAN_VARIANT", "cn-personal")
    if variant_id not in VARIANTS:
        variant_id = "cn-personal"

    snapshot = {"schema": 1,
                "source": {"cookie": "console-cookie", "apikey": "apikey-probe"}.get(mode, mode),
                "plan": "personal", "probed_at": now_iso(),
                "ok": False, "stale": False, "last_error": None,
                "error_kind": None, "plan_level": None, "windows": [],
                "credits": None, "diagnostics": None, "reset_cards": [],
                "subscription": None}
    last = read_json(state_path) or {}
    try:
        if mode == "apikey":
            windows, plan, diag = probe_apikey_mode()
            reset_cards = []
            subscription = None
        else:
            windows, plan, diag, reset_cards, subscription = probe_cookie_mode(variant_id)
        flags = zero_without_rollover(windows, last.get("windows"))
        if flags:
            diag["zero_without_rollover"] = flags
        snapshot.update({"ok": True, "plan_level": plan, "windows": windows,
                         "diagnostics": diag, "reset_cards": reset_cards,
                         "subscription": subscription})
        # credits 只报 7 天池（两个窗口是两个池子，相加无意义）
        wk = next((w for w in windows if w["id"] == "weekly"), None)
        if wk and wk.get("total"):
            snapshot["credits"] = {"total": wk["total"],
                                   "used": round(wk["total"] - (wk["remaining"] or 0), 1),
                                   "remaining": wk["remaining"], "unit": "Credits"}
        atomic_write(state_path, snapshot)
    except ProbeError as e:
        snapshot["last_error"], snapshot["error_kind"] = str(e), e.kind
        if e.diag:
            snapshot["diagnostics"] = e.diag
    except Exception as e:  # 网络抖动等：也走 stale 回填，不让前端白屏
        snapshot["last_error"], snapshot["error_kind"] = f"{type(e).__name__}: {e}", "unavailable"

    if not snapshot["ok"] and last.get("windows"):
        # 失败保留上次成功值（与 Sub2API CN 探测同策略），显式标 stale
        snapshot["windows"] = last["windows"]
        snapshot["plan_level"] = last.get("plan_level") or snapshot["plan_level"]
        snapshot["credits"] = last.get("credits")
        snapshot["reset_cards"] = last.get("reset_cards") or []
        snapshot["subscription"] = last.get("subscription")
        snapshot["stale"] = True
        snapshot["last_ok_at"] = last.get("probed_at")

    try:
        os.makedirs(os.path.dirname(os.path.abspath(log_path)), exist_ok=True)
        with open(log_path, "a", encoding="utf-8") as f:
            f.write(json.dumps({"ts": snapshot["probed_at"], "mode": mode,
                                "ok": snapshot["ok"], "kind": snapshot["error_kind"],
                                "err": (snapshot["last_error"] or "")[:200]},
                               ensure_ascii=False) + "\n")
    except Exception:
        pass

    atomic_write(out, snapshot)
    print(json.dumps({"ok": snapshot["ok"], "stale": snapshot["stale"],
                      "windows": len(snapshot["windows"]),
                      "error": snapshot["last_error"]}, ensure_ascii=False))
    return 0 if (snapshot["ok"] or snapshot["stale"]) else 1


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