# -*- coding: utf-8 -*-
"""
fix-node-model.py —— Hermes 节点模型链路抢修（Windows / Linux / macOS 通用）

背景（2026-09-30）：上游面板 gemini 全系配额耗尽（"Individual quota reached"）、
doubao 月配额到 10-09；同时节点本地 v2ray/xray 的全局代理会把「模型请求」也代理掉，
那条出口对面板返 502。两者叠加的表现是：一直报错 / 每个候选都失败。

本脚本做四件事（都可单独关闭）：
  1. 备份 config.yaml（config-backups/，只留最近 2 份）
  2. 主模型 / delegation 切到「实测可用」的模型（默认 glm-5.3 → deepseek-v4-flash）
  3. 备用链重排：实测可用的排前面，gemini/doubao 这类已知死的挪到链尾（条目原样保留，不动 api_key）
  4. 代理：给 .env 和启动脚本补 NO_PROXY（让模型走直连，TG 仍走代理）
  5. 摘掉与主模型冲突的会话级 /model 覆盖（sessions.json 的 model_override）

用法：
  python fix-node-model.py --check          # 只读：把当前状态 + 实测结果全部打印出来
  python fix-node-model.py                  # 执行修复
  python fix-node-model.py --revert         # 从最近一次备份还原
  python fix-node-model.py --home "D:\\Program Files\\hermes"
开关：--no-proxy（不动代理） --no-chain（不重排链） --no-pins（不摘会话覆盖）
"""
import argparse
import glob
import json
import os
import shutil
import socket
import ssl
import sys
import time

try:
    import yaml
except ImportError:
    print("ERROR: 缺 pyyaml。请用她自带的 python 跑，例如：")
    print('  "D:\\Program Files\\hermes\\.venv\\Scripts\\python.exe" fix-node-model.py --check')
    sys.exit(1)

CAND_MODELS = ["glm-5.3", "glm-5.2", "deepseek-v4-flash", "agnes-2.5-pro"]
DEAD_HINTS = ("gemini", "doubao")          # 已知配额耗尽/被风控的族
PANEL_HOSTS = "cc.3602020.xyz,api.office3.pp.ua,.3602020.xyz,.office3.pp.ua"
NO_PROXY_BASE = "localhost,127.0.0.1,::1,10.0.0.0/8,172.16.0.0/12,192.168.0.0/16"


def find_home(explicit=None):
    cands = [explicit, os.environ.get("HERMES_HOME"),
             os.path.join(os.environ.get("USERPROFILE", ""), ".hermes"),
             os.path.expanduser("~/.hermes"),
             r"D:\Program Files\hermes", "D:\\soft\\hermes", "C:\\hermes"]
    for c in cands:
        if c and os.path.isfile(os.path.join(c, "config.yaml")):
            return c
    return None


def load_yaml(p):
    with open(p, encoding="utf-8") as f:
        return yaml.safe_load(f) or {}


def dump_yaml(obj, p):
    tmp = p + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        yaml.safe_dump(obj, f, allow_unicode=True, sort_keys=False, width=10 ** 6)
    os.replace(tmp, p)


def http_status(url, key, model, proxy=None, timeout=12):
    """一次 chat 级硬超时实测；返回 (ok, 说明)"""
    import urllib.request
    body = json.dumps({"model": model, "messages": [{"role": "user", "content": "hi"}],
                       "max_tokens": 4}).encode()
    req = urllib.request.Request(url.rstrip("/") + "/chat/completions", data=body,
        headers={"Content-Type": "application/json", "Authorization": "Bearer " + key})
    handlers = [urllib.request.ProxyHandler({"http": proxy, "https": proxy})] if proxy \
        else [urllib.request.ProxyHandler({})]
    opener = urllib.request.build_opener(*handlers)
    t0 = time.time()
    try:
        opener.open(req, timeout=timeout).read()
        return True, "HTTP 200 (%.1fs)" % (time.time() - t0)
    except urllib.error.HTTPError as e:
        txt = ""
        try:
            txt = json.loads(e.read().decode()).get("error", {}).get("message", "")[:50]
        except Exception:
            pass
        return False, "HTTP %s %s" % (e.code, txt)
    except Exception as e:
        return False, "%s: %s" % (type(e).__name__, str(e)[:50])


def endpoints(cfg):
    """从 config 里取入口：优先主模型那条，其次 providers"""
    out = []
    m = cfg.get("model") or {}
    if m.get("base_url") and m.get("api_key"):
        out.append((m.get("provider") or "main", m["base_url"], m["api_key"]))
    for e in (cfg.get("fallback_providers") or []):
        base, key = e.get("base_url"), e.get("api_key")
        if base and key and all(base != b for _, b, _ in out):
            out.append((e.get("provider"), base, key))
    for name, p in (cfg.get("providers") or {}).items():
        if isinstance(p, dict) and p.get("api") and p.get("api_key"):
            if all(p["api"] != b for _, b, _ in out):
                out.append((name, p["api"], p["api_key"]))
    return out


def proxy_env_of(cfg_dir):
    """读 .env 里的代理设置"""
    p = os.path.join(cfg_dir, ".env")
    got = {}
    if os.path.isfile(p):
        for ln in open(p, encoding="utf-8", errors="ignore"):
            ln = ln.strip()
            if not ln or ln.startswith("#") or "=" not in ln:
                continue
            k, v = ln.split("=", 1)
            got[k.strip()] = v.strip()
    return got


def launcher_files(cfg_dir, extra_dirs=()):
    hits = []
    dirs = [cfg_dir] + [d for d in extra_dirs if d]
    for d in dirs:
        for pat in ("*.bat", "*.cmd", "*.vbs", "*.ps1", "*.sh"):
            for f in glob.glob(os.path.join(d, pat)):
                try:
                    txt = open(f, encoding="utf-8", errors="ignore").read()
                except OSError:
                    continue
                if "PROXY" in txt.upper():
                    hits.append(f)
    return hits


def patch_proxy(cfg_dir, dry, extra_dirs=()):
    """给 .env 与启动脚本补 NO_PROXY（保留各自原有变量）"""
    want = NO_PROXY_BASE + "," + PANEL_HOSTS
    changed = []

    # 1) .env
    p = os.path.join(cfg_dir, ".env")
    if os.path.isfile(p):
        lines = open(p, encoding="utf-8", errors="ignore").read().splitlines()
        out, has = [], False
        for ln in lines:
            if ln.strip().upper().startswith("NO_PROXY="):
                has = True
                cur = ln.split("=", 1)[1]
                if all(h in cur for h in PANEL_HOSTS.split(",")):
                    out.append(ln)
                else:
                    out.append("NO_PROXY=" + want)
                    changed.append(".env: NO_PROXY 补上面板域名")
            else:
                out.append(ln)
        if not has:
            out.append("NO_PROXY=" + want)
            changed.append(".env: 新增 NO_PROXY")
        if not dry and changed:
            shutil.copy2(p, p + ".bak-%s" % time.strftime("%Y%m%d-%H%M%S"))
            open(p, "w", encoding="utf-8").write("\n".join(out) + "\n")

    # 2) 启动脚本（.bat/.ps1/.vbs 等）—— 程序目录里的启动脚本不在主目录，需另传 --launcher-dir
    for f in launcher_files(cfg_dir, extra_dirs):
        txt = open(f, encoding="utf-8", errors="ignore").read()
        if PANEL_HOSTS.split(",")[0] in txt:
            continue
        ext = os.path.splitext(f)[1].lower()
        if ext in (".bat", ".cmd"):
            newline = 'set "NO_PROXY=%s"\r\n' % want
        elif ext == ".ps1":
            newline = '$env:NO_PROXY="%s"\r\n' % want
        else:                                     # .vbs/.sh 之类的先不硬塞
            changed.append("%s: 含代理变量但类型未知，未自动改（请人工看一眼）" % os.path.basename(f))
            continue
        if dry:
            changed.append("%s: 会插入 NO_PROXY" % os.path.basename(f))
            continue
        shutil.copy2(f, f + ".bak-%s" % time.strftime("%Y%m%d-%H%M%S"))
        # 插在最后一条代理 set 之后（保持 @echo off/chcp 在最前），没有就插到最前
        lines = txt.splitlines(True)
        idx = 0
        for i, ln in enumerate(lines):
            if "PROXY" in ln.upper() and "NO_PROXY" not in ln.upper():
                idx = i + 1
        lines.insert(idx, newline)
        open(f, "w", encoding="utf-8").write("".join(lines))
        changed.append("%s: 已插入 NO_PROXY" % os.path.basename(f))
    return changed


def clear_pins(cfg_dir, target, dry):
    """摘掉与目标主模型冲突的会话级 /model 覆盖"""
    p = os.path.join(cfg_dir, "sessions", "sessions.json")
    if not os.path.isfile(p):
        return []
    try:
        data = json.load(open(p, encoding="utf-8"))
    except Exception as e:
        return ["sessions.json 读不了: %s" % e]
    hits = []
    if isinstance(data, dict):
        for k, v in data.items():
            if isinstance(v, dict) and v.get("model_override") and v["model_override"] != target:
                hits.append("%s -> %s" % (k, v["model_override"]))
                if not dry:
                    v.pop("model_override", None)
    if hits and not dry:
        shutil.copy2(p, p + ".bak-%s" % time.strftime("%Y%m%d-%H%M%S"))
        json.dump(data, open(p, "w", encoding="utf-8"), ensure_ascii=False, indent=1)
    return hits


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--home")
    ap.add_argument("--check", action="store_true")
    ap.add_argument("--revert", action="store_true")
    ap.add_argument("--no-proxy", action="store_true")
    ap.add_argument("--no-chain", action="store_true")
    ap.add_argument("--no-pins", action="store_true")
    ap.add_argument("--launcher-dir", action="append", default=[],
                    help="额外扫描启动脚本的目录（Windows：程序目录，如 \"D:\\Program Files\\hermes\"）")
    a = ap.parse_args()

    home = find_home(a.home)
    if not home:
        print("ERROR: 找不到 config.yaml，请用 --home 指定 Hermes 主目录")
        return 1
    cfg_path = os.path.join(home, "config.yaml")
    cfg = load_yaml(cfg_path)
    m = cfg.get("model") or {}
    dele = cfg.get("delegation") or {}
    chain = cfg.get("fallback_providers") or []

    print("== Hermes 主目录:", home)
    print("== 主模型: %s @ %s | base_url: %s" % (m.get("default"), m.get("provider"), m.get("base_url")))
    print("== 子模型:", dele.get("model"))
    print("== 备用链 %d 条:" % len(chain))
    for i, e in enumerate(chain):
        print("   %2d. %s/%s" % (i + 1, e.get("provider"), e.get("model")))

    if a.revert:
        baks = sorted(glob.glob(cfg_path + ".bak-*") + glob.glob(os.path.join(home, "config-backups", "*")))
        if not baks:
            print("没有可还原的备份"); return 1
        shutil.copy2(baks[-1], cfg_path)
        print("✅ 已还原:", baks[-1])
        return 0

    # ---- 实测：直连 vs 走代理 ----
    eps = endpoints(cfg)
    penv = proxy_env_of(home)
    proxy = penv.get("HTTPS_PROXY") or penv.get("HTTP_PROXY") or os.environ.get("HTTPS_PROXY")
    print("\n== 代理设置: HTTPS_PROXY=%s | NO_PROXY=%s"
          % (penv.get("HTTPS_PROXY"), (penv.get("NO_PROXY") or "(无)")))
    print("== 启动脚本含代理变量:", [os.path.basename(f) for f in launcher_files(home, a.launcher_dir)] or "无")

    results, best = {}, None
    for model in CAND_MODELS:
        info = []
        ok_any = False
        for name, base, key in eps[:2]:
            ok, why = http_status(base, key, model, proxy=None)
            info.append("直连%s %s" % (name, "✅" if ok else "❌" + why))
            if ok:
                ok_any = True
            if proxy:
                ok2, why2 = http_status(base, key, model, proxy=proxy)
                info.append("代理%s %s" % (name, "✅" if ok2 else "❌" + why2))
        results[model] = (ok_any, " | ".join(info))
        print("== %-18s %s" % (model, results[model][1]))
        if ok_any and best is None:
            best = model
    alives = [k for k, v in results.items() if v[0]]
    print("== 实测可用:", alives or "(全挂)")
    if best is None:
        best = m.get("default")

    if a.check:
        print("\n== [--check] 将要做的改动：")
        print("   主模型/子模型 →", best)
        if not a.no_chain:
            alive_c = [e for e in chain if str(e.get("model")) not in DEAD_HINTS and
                       (e.get("model") in alives or not alives)]
            dead_c = [e for e in chain if e not in alive_c]
            print("   链重排: 可用 %d 条提前, 移尾 %d 条 (%s)"
                  % (len(alive_c), len(dead_c),
                     ", ".join(sorted({str(e.get("model")) for e in dead_c})) or "无"))
        if not a.no_proxy:
            print("   代理: 给 .env/启动脚本补 NO_PROXY（面板域名直连）")
        if not a.no_pins:
            print("   会话覆盖: 摘掉与 %s 冲突的 model_override" % best)
        return 0

    # ---- 执行 ----
    bdir = os.path.join(home, "config-backups")
    os.makedirs(bdir, exist_ok=True)
    shutil.copy2(cfg_path, os.path.join(bdir, "config-before-nodefix-%s.yaml"
                                       % time.strftime("%Y%m%d-%H%M%S")))
    baks = sorted(glob.glob(os.path.join(bdir, "config-before-nodefix-*.yaml")))
    for old in baks[:-2]:
        try: os.remove(old)
        except OSError: pass
    print("\n== 📦 备份到", bdir)

    changed = []
    if m.get("default") != best:
        changed.append("model.default: %s → %s" % (m.get("default"), best))
        m["default"] = best
    if dele.get("model") and dele.get("model") != best:
        changed.append("delegation.model: %s → %s" % (dele.get("model"), best))
        dele["model"] = best
    cfg["model"], cfg["delegation"] = m, dele

    if not a.no_chain and chain:
        alive_c = [e for e in chain if str(e.get("model")) not in DEAD_HINTS and
                   (e.get("model") in alives or not alives)]
        dead_c = [e for e in chain if e not in alive_c]
        if alive_c + dead_c != chain:
            cfg["fallback_providers"] = alive_c + dead_c
            changed.append("备用链: %d 条移尾 (%s)"
                           % (len(dead_c), ", ".join(sorted({str(e.get("model")) for e in dead_c}))))
    dump_yaml(cfg, cfg_path)

    for c in changed:
        print("   ✔", c)

    if not a.no_proxy:
        for c in patch_proxy(home, dry=False, extra_dirs=a.launcher_dir):
            print("   ✔ 代理:", c)
    if not a.no_pins:
        pins = clear_pins(home, best, dry=False)
        for p in pins:
            print("   ✔ 会话覆盖已摘:", p)

    chk = load_yaml(cfg_path)
    print("\n== ✅ 验收: 主模型 %s | 子模型 %s | 链前 3 条: %s"
          % ((chk.get("model") or {}).get("default"),
             (chk.get("delegation") or {}).get("model"),
             ", ".join("%s/%s" % (e.get("provider"), e.get("model"))
                       for e in (chk.get("fallback_providers") or [])[:3])))
    print("== 👉 下一条新消息自动生效（改的是 config，无覆盖的会话下一轮就读新值）")
    print("== 👉 若改了代理/启动脚本，需要重开那个 launcher 才生效（或下次开机）")
    return 0


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