#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
agent.py — GPU 节点端代理：**出站**连接门户（WebSocket），注册/心跳/隧道/工作区同步。

为什么有它：GPU 机器不再需要任何入站端口、隧道、证书（§0.3 的 Meshnet 限制、Vast 的端口映射
坑全部消失）。门户把网页/账号/信令收走，节点只要能上网就能被分配给任何客户端；媒体仍是
浏览器 ↔ 节点直连优先，打不通走门户的 TURN 兜底。

配置（环境变量优先，其次同目录 agent_config.json）：
  PORTAL_URL          门户地址，如 https://fs.hawkmind.ai
  PORTAL_NODE_TOKEN   节点 token。**可以不填**：没有令牌时 agent 会用机器指纹向门户 /api/node/enroll 自动注册并把拿到的
                      令牌写回 agent_config.json（2026-09-06 user 拍板：装完即加入系统、后台默认启用、不认识的机器管理员拉黑）
  NODE_TOKEN          本机 server.py 的 X-Token（与 run_node.cmd / onstart 里的一致，默认 dev）
  NODE_HTTP           本机 server.py 地址（默认 https://127.0.0.1:8080，自签证书不校验；没证书就 http）
  MANAGE_SERVER       1=由 agent 拉起并看护 server.py（Vast 用）；0=server.py 由计划任务/其他方式常驻（ai002 用）
  NODE_NAME           可选，注册时的名字提示

🔴 只用 aiohttp（节点已装），不引入新依赖；不 import torch（重）。GPU 信息用 nvidia-smi。
🔴 日志同时落 agent.log：Vast 上日志 API 是唯一回传通道，本文件 print 会进容器日志。
"""
import asyncio
import base64
import hashlib
import json
import os
import platform
import shutil
import socket
import ssl
import subprocess
import sys
import time
import uuid

import aiohttp

HERE = os.path.dirname(os.path.abspath(__file__))
ASSETS = os.path.join(HERE, "assets")
CFG_PATH = os.path.join(HERE, "node_config.json")
AGENT_CFG_PATH = os.path.join(HERE, "agent_config.json")
AGENT_VERSION = "1.1.0"


def _load_cfg():
    cfg = {}
    p = os.path.join(HERE, "agent_config.json")
    if os.path.exists(p):
        try:
            cfg = json.load(open(p, encoding="utf-8-sig"))     # utf-8-sig：记事本/PowerShell 写的带 BOM 文件也能读（2026-09-06 实遇）
        except Exception as e:
            print("agent_config.json 读取失败:", e, flush=True)
    for k in ("PORTAL_URL", "PORTAL_NODE_TOKEN", "NODE_TOKEN", "NODE_HTTP", "MANAGE_SERVER", "NODE_NAME", "NODE_MODE"):
        if os.environ.get(k):
            cfg[k] = os.environ[k]
    cfg.setdefault("NODE_TOKEN", "dev")
    cfg.setdefault("NODE_HTTP", "https://127.0.0.1:8080")
    cfg.setdefault("MANAGE_SERVER", "0")
    cfg.setdefault("NODE_MODE", "pipeline")
    return cfg


CFG = _load_cfg()
_LOGF = os.path.join(HERE, "agent.log")


def log(*a):
    s = "[%s] %s" % (time.strftime("%H:%M:%S"), " ".join(str(x) for x in a))
    print(s, flush=True)
    try:
        with open(_LOGF, "a", encoding="utf-8") as f:
            f.write(s + "\n")
    except Exception:
        pass


NOVERIFY = ssl.create_default_context()
NOVERIFY.check_hostname = False
NOVERIFY.verify_mode = ssl.CERT_NONE


def machine_fingerprint():
    """机器指纹：Windows 用注册表 MachineGuid，Linux 用 /etc/machine-id，都没有就退到网卡 MAC。sha256 取前 32 位。
    只用来让门户认出「同一台机器」（重装/丢配置不重复建节点、拉黑后换个目录装也认得出、桌面版判「本机 GPU」），不含隐私。
    🔴 desktop/localsrv.py 里有一份同算法的副本（桌面版用它判本机），改这里要一起改。"""
    raw = ""
    try:
        if os.name == "nt":
            import winreg
            k = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Microsoft\Cryptography", 0,
                               winreg.KEY_READ | winreg.KEY_WOW64_64KEY)
            raw = str(winreg.QueryValueEx(k, "MachineGuid")[0])
        else:
            for p in ("/etc/machine-id", "/var/lib/dbus/machine-id"):
                if os.path.exists(p):
                    raw = open(p).read().strip()
                    break
    except Exception:
        raw = ""
    if not raw:
        raw = "mac:%012x" % uuid.getnode()
    return hashlib.sha256(("faceswap-node:" + raw).encode("utf-8")).hexdigest()[:32]


def instance_key(port):
    """同一台机器 + 同一端口 = 同一个安装（重装不重复建节点）；同一台机器第二个端口 = 第二个节点。"""
    return hashlib.sha256((machine_fingerprint() + ":" + str(port)).encode("utf-8")).hexdigest()[:32]


def lan_ips():
    """本机各网卡的私网 IPv4（门户拿它判客户端与节点是否在同一局域网）。"""
    out = []
    try:
        for info in socket.getaddrinfo(socket.gethostname(), None, socket.AF_INET):
            ip = info[4][0]
            p = ip.split(".")
            a, b = int(p[0]), int(p[1])
            if (a == 10 or (a == 192 and b == 168) or (a == 172 and 16 <= b <= 31)) and ip not in out:
                out.append(ip)
    except Exception:
        pass
    return out[:16]


def save_agent_token(token):
    """把自动注册拿到的令牌写回 agent_config.json（下次启动直接用，不再重复注册）。"""
    cfg = {}
    if os.path.exists(AGENT_CFG_PATH):
        try:
            cfg = json.load(open(AGENT_CFG_PATH, encoding="utf-8-sig"))
        except Exception:
            cfg = {}
    cfg["PORTAL_NODE_TOKEN"] = token
    cfg.setdefault("PORTAL_URL", CFG.get("PORTAL_URL", ""))
    tmp = AGENT_CFG_PATH + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(cfg, f, ensure_ascii=False, indent=2)
    os.replace(tmp, AGENT_CFG_PATH)


def gpu_info():
    """(名字, 显存MB, 利用率%, 已用显存MB)；没有 nvidia-smi 就全空。"""
    exe = shutil.which("nvidia-smi") or (r"C:\Windows\System32\nvidia-smi.exe" if os.name == "nt" else "nvidia-smi")
    try:
        out = subprocess.run([exe, "--query-gpu=name,memory.total,utilization.gpu,memory.used",
                              "--format=csv,noheader,nounits"], capture_output=True, text=True, timeout=8).stdout
        parts = [p.strip() for p in out.strip().splitlines()[0].split(",")]
        return parts[0], int(float(parts[1])), int(float(parts[2])), int(float(parts[3]))
    except Exception:
        return "", 0, None, None


def read_node_config():
    try:
        return json.load(open(CFG_PATH, encoding="utf-8"))
    except Exception:
        return {}


def write_node_config(patch: dict, drop=()):
    cfg = read_node_config()
    for k in drop:
        cfg.pop(k, None)
    cfg.update(patch)
    tmp = CFG_PATH + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(cfg, f, ensure_ascii=False, indent=2)
    os.replace(tmp, CFG_PATH)


def _labels(path):
    try:
        return json.load(open(path, encoding="utf-8"))
    except Exception:
        return {}


def _write_labels(path, d):
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(d, f, ensure_ascii=False, indent=1)
    os.replace(tmp, path)


PERSONA_META = os.path.join(ASSETS, "personas.json")
BG_META = os.path.join(ASSETS, "backgrounds.json")
WS_STATE = os.path.join(ASSETS, "_workspace.json")   # 记账：本会话为哪个用户下载/新生成了哪些文件
SYS_STATE = os.path.join(ASSETS, "_system.json")     # 记账：门户下发的系统素材（跨会话保留，下线时按它清理）


def _file_sha(p):
    """与门户 workspace.file_sha 同口径（sha256 前 32 位）；文件不存在返回空串。"""
    import hashlib
    try:
        h = hashlib.sha256()
        with open(p, "rb") as f:
            for chunk in iter(lambda: f.read(1 << 20), b""):
                h.update(chunk)
        return h.hexdigest()[:32]
    except Exception:
        return ""


def _is_asset(name):
    low = name.lower()
    return (low.startswith("persona") or low.startswith("bg") or low.startswith("bgh__")) and \
        low.endswith((".jpg", ".jpeg", ".png", ".webp"))


class Agent:
    def __init__(self):
        self.portal = CFG.get("PORTAL_URL", "").rstrip("/")
        self.token = CFG.get("PORTAL_NODE_TOKEN", "")
        self.token_from_env = bool(os.environ.get("PORTAL_NODE_TOKEN"))   # Vast 走 env 给令牌：被停用时不该自己再去注册
        self.http_token = CFG.get("NODE_TOKEN", "dev")
        self.local = CFG.get("NODE_HTTP", "https://127.0.0.1:8080").rstrip("/")
        try:
            self.local_port = int(self.local.rsplit(":", 1)[1])
        except Exception:
            self.local_port = 0
        self.manage = str(CFG.get("MANAGE_SERVER", "0")) in ("1", "true", "yes")
        self.ws = None
        self.session = None            # {"sid","user_id"}
        self.server_proc = None
        self.sess_http = None

    # ---------------- 本机 server.py
    async def local_http(self, method, path, query="", headers=None, body=b"", timeout=120):
        url = self.local + path + (("?" + query) if query else "")
        h = dict(headers or {})
        h["X-Token"] = self.http_token
        async with self.sess_http.request(method, url, data=body if body else None, headers=h,
                                          ssl=NOVERIFY if url.startswith("https") else None,
                                          timeout=aiohttp.ClientTimeout(total=timeout)) as r:
            data = await r.read()
            return r.status, {k: v for k, v in r.headers.items()
                              if k.lower() in ("content-type", "cache-control", "content-disposition")}, data

    async def local_json(self, method, path, obj=None, timeout=60):
        body = json.dumps(obj).encode() if obj is not None else b""
        st, _, data = await self.local_http(method, path, headers={"Content-Type": "application/json"} if obj is not None else None,
                                            body=body, timeout=timeout)
        try:
            return st, json.loads(data.decode("utf-8") or "null")
        except Exception:
            return st, None

    async def wait_local_health(self, timeout=1200):
        t0 = time.time()
        last_tail = 0
        while time.time() - t0 < timeout:
            try:
                st, h = await self.local_json("GET", "/health", timeout=8)
                if st == 200 and h:
                    return h
            except Exception:
                pass
            if self.manage and self.server_proc is not None and self.server_proc.poll() is not None:
                log("server.py 退出 rc=%s，重新拉起" % self.server_proc.returncode)
                self.start_server()
            # 🔴 Vast 上只有容器 stdout 能看到（日志 API），server.log 在文件里看不见：等待期间每分钟把它末尾两行抄进来，
            #    否则"server.py 20 分钟没就绪"完全无从诊断（2026-09-06 租赁验证时实遇：只看到一行"等待本机 server.py 就绪"）
            if time.time() - last_tail > 60:
                last_tail = time.time()
                try:
                    with open(os.path.join(HERE, "server.log"), "rb") as f:
                        f.seek(0, 2)
                        f.seek(max(0, f.tell() - 600))
                        lines = [x for x in f.read().decode("utf-8", "replace").splitlines() if x.strip()][-2:]
                    log("仍在等 server.py（%ds）… server.log 末尾: %s" % (int(time.time() - t0), " | ".join(x[:160] for x in lines)))
                except Exception:
                    log("仍在等 server.py（%ds）… （server.log 还没生成）" % int(time.time() - t0))
            await asyncio.sleep(3)
        return None

    def start_server(self):
        env = dict(os.environ)
        env["NODE_TOKEN"] = self.http_token
        env.setdefault("NODE_MODE", CFG.get("NODE_MODE", "pipeline"))
        env.setdefault("PYTHONUNBUFFERED", "1")
        with open(os.path.join(HERE, ".mode"), "w") as f:
            f.write(env["NODE_MODE"])
        logf = open(os.path.join(HERE, "server.log"), "a", encoding="utf-8")
        self.server_proc = subprocess.Popen([sys.executable, "-u", "server.py", "--port", "8080"], cwd=HERE,
                                            env=env, stdout=logf, stderr=subprocess.STDOUT)
        log("已拉起 server.py pid=%d" % self.server_proc.pid)
        if not os.path.exists(os.path.join(HERE, "cert.pem")):
            self.local = "http://127.0.0.1:8080"

    # ---------------- 工作区
    async def ws_download(self, url, dest):
        async with self.sess_http.get(url, headers={"Authorization": "Bearer " + self.token},
                                      timeout=aiohttp.ClientTimeout(total=120)) as r:
            if r.status != 200:
                raise RuntimeError("HTTP %d" % r.status)
            data = await r.read()
        with open(dest + ".tmp", "wb") as f:
            f.write(data)
        os.replace(dest + ".tmp", dest)
        return len(data)

    async def session_start(self, msg):
        sid, uid = msg.get("sid"), msg.get("user_id")
        ws = msg.get("workspace") or {}
        # ① 会话级 ICE（含 TURN 时限凭据）→ node_config.json，server.py 的 /offer 现读
        write_node_config({"ice_servers": msg.get("ice_servers") or []})
        # ② key（bg_harmonize 用）；不存在的属性 /tune 会跳过
        keys = msg.get("keys") or {}
        if keys:
            try:
                await self.local_json("POST", "/tune", {k: v for k, v in keys.items() if v})
            except Exception as e:
                log("灌 key 失败:", e)
        # ③ 先记账再下载：万一中途断电，下次会话开始也知道该清什么
        state = _labels(WS_STATE)
        if state.get("sid") and state.get("sid") != sid:
            await self.session_end({"sid": state["sid"], "user_id": state.get("user_id"), "reason": "stale"})
        before = set(f for f in os.listdir(ASSETS) if _is_asset(f))
        state = {"sid": sid, "user_id": uid, "downloaded": [], "before": sorted(before), "t": time.time()}
        _write_labels(WS_STATE, state)
        plabels, blabels = _labels(PERSONA_META), _labels(BG_META)
        n = skipped = 0
        sys_now = {}                                    # 本次清单里的系统素材：name -> sha（用于下线清理）
        for f in ws.get("files") or []:
            name = os.path.basename(str(f.get("name") or ""))
            if not _is_asset(name):
                continue
            dst = os.path.join(ASSETS, name)
            # 🔴 系统素材（所有账号共享的预置人设等）：会话结束不删（见 session_end），所以要判断本地这份是不是
            #    还是门户上那一份——**按内容哈希比，不能只看大小**：换一张同尺寸的图会永远同步不过去。
            sys_asset = bool(f.get("system"))
            sha = str(f.get("sha") or "")
            if sys_asset:
                sys_now[name] = sha
            if sys_asset and sha and _file_sha(dst) == sha:
                (plabels if f.get("kind") == "persona" else blabels)[name] = f.get("label") or name
                skipped += 1
                continue
            try:
                await self.ws_download(f["url"], dst)
                if not sys_asset:
                    state["downloaded"].append(name)      # 只有用户私有文件登记进"会话结束要清掉"的清单
                (plabels if f.get("kind") == "persona" else blabels)[name] = f.get("label") or name
                n += 1
            except Exception as e:
                log("工作区文件下载失败 %s: %s" % (name, e))
        # 系统素材下线清理：上次装过、这次清单里没有的，删掉——否则节点上会越堆越多没人认领的旧人设
        old_sys = _labels(SYS_STATE)
        for name in list(old_sys.keys()):
            if name in sys_now or not _is_asset(name):
                continue
            try:
                os.remove(os.path.join(ASSETS, name))
                log("系统素材已下线，删除本地副本: %s" % name)
            except FileNotFoundError:
                pass
            except Exception as e:
                log("删除下线素材 %s 失败: %s" % (name, e))
            plabels.pop(name, None)
            blabels.pop(name, None)
        _write_labels(SYS_STATE, sys_now)
        _write_labels(PERSONA_META, plabels)
        _write_labels(BG_META, blabels)
        _write_labels(WS_STATE, state)
        # ④ 预设（server.py 每次请求现读 presets.json，写文件即生效）
        # 🔴 节点本地的 presets.json 是**直连模式用户的数据**（ai002 上 user 用 v1.3.0 存的预设就在这里）：
        #    会话开始前先备份到 presets.local.json，会话结束恢复；门户账号的预设只在会话期间生效。
        #    2026-09-06 第一版没备份，把 user 在 ai002 上的预设覆盖成了空——这一段就是为此而写。
        pres = ws.get("presets")
        if isinstance(pres, dict):
            pp = os.path.join(HERE, "presets.json")
            bak = os.path.join(HERE, "presets.local.json")
            if os.path.exists(pp) and not os.path.exists(bak):
                try:
                    shutil.copy2(pp, bak)
                except Exception as e:
                    log("备份本地预设失败:", e)
            _write_labels(pp, pres)
        # ⑤ 隐私：管线里的人设/背景回到默认，不把上一位用户的脸留给下一位
        try:
            await self.local_json("POST", "/tune", {"persona_name": "persona.png", "bg_name": "background.jpg"})
        except Exception as e:
            log("重置人设/背景失败:", e)
        self.session = {"sid": sid, "user_id": uid}
        log("会话开始 %s 用户 %s：同步 %d 个素材" % (sid, uid, n))

    async def session_end(self, msg):
        sid, uid = msg.get("sid"), msg.get("user_id")
        state = _labels(WS_STATE)
        if not state:
            self.session = None
            return
        uid = uid or state.get("user_id")
        before = set(state.get("before") or [])
        now_files = set(f for f in os.listdir(ASSETS) if _is_asset(f))
        # 🔴 系统素材要从"会话期间新增"里排除：它们是门户下发的公共素材，既不该回传进某个用户的私有工作区，
        #    也不该在会话结束时删掉（删了下次每台节点都要重下）。
        sys_assets = set(_labels(SYS_STATE).keys())
        new = sorted(now_files - before - sys_assets)           # 会话期间新生成/上传的（不含系统素材）
        plabels, blabels = _labels(PERSONA_META), _labels(BG_META)
        # ① 新文件传回门户（上传时隧道已镜像过，这里补 AI 候选/光照变体这类节点自己生成的）
        for name in new:
            p = os.path.join(ASSETS, name)
            try:
                label = plabels.get(name) or blabels.get(name) or ""
                async with self.sess_http.post("%s/api/node/upload/%s/%s" % (self.portal, uid, name),
                                               params={"label": label}, data=open(p, "rb").read(),
                                               headers={"Authorization": "Bearer " + self.token},
                                               timeout=aiohttp.ClientTimeout(total=120)) as r:
                    if r.status != 200:
                        log("回传 %s 失败 HTTP %d" % (name, r.status))
            except Exception as e:
                log("回传 %s 失败: %s" % (name, e))
        # ② 管线回默认，再删文件（顺序不能反：删了正在用的图管线会报错）
        try:
            await self.local_json("POST", "/tune", {"persona_name": "persona.png", "bg_name": "background.jpg"})
        except Exception:
            pass
        removed = 0
        for name in set(state.get("downloaded") or []) | set(new):
            try:
                os.remove(os.path.join(ASSETS, name)); removed += 1
            except FileNotFoundError:
                pass
            except Exception as e:
                log("删除 %s 失败: %s" % (name, e))
            plabels.pop(name, None); blabels.pop(name, None)
        _write_labels(PERSONA_META, plabels)
        _write_labels(BG_META, blabels)
        pp, bak = os.path.join(HERE, "presets.json"), os.path.join(HERE, "presets.local.json")
        if os.path.exists(bak):                     # 恢复直连模式用户的本地预设（见 session_start ④）
            try:
                os.replace(bak, pp)
            except Exception as e:
                log("恢复本地预设失败:", e)
        else:
            _write_labels(pp, {"presets": [], "current": None})
        write_node_config({}, drop=("ice_servers",))
        try:
            os.remove(WS_STATE)
        except FileNotFoundError:
            pass
        # 远程协助旁观流（§15.11）：会话结束后旁观者没有东西可看，关掉；旧版 server.py 没这个端点会 404，只记日志
        try:
            st, _r = await self.local_json("POST", "/watch/stop", {})
            if st != 200:
                log("关闭旁观连接返回 HTTP %s" % st)
        except Exception as e:
            log("关闭旁观连接失败:", e)
        self.session = None
        log("会话结束 %s：回传 %d 个新文件，清理 %d 个文件（%s）" % (sid, len(new), removed, msg.get("reason", "")))

    # ---------------- 门户消息
    async def handle(self, msg):
        t = msg.get("type")
        if t == "http":
            rid = msg.get("id")
            try:
                body = base64.b64decode(msg.get("body_b64") or "")
                st, hd, data = await self.local_http(msg.get("method", "GET"), msg.get("path", "/"), msg.get("query", ""),
                                                     msg.get("headers") or {}, body,
                                                     timeout=170 if msg.get("path") in ("/offer", "/watch/offer") else 120)
                await self.send({"type": "http_resp", "id": rid, "status": st, "headers": hd,
                                 "body_b64": base64.b64encode(data).decode("ascii")})
            except Exception as e:
                await self.send({"type": "http_resp", "id": rid, "error": "%s: %s" % (type(e).__name__, e)})
        elif t == "config":
            patch = {}
            if msg.get("public_host"):
                patch["public_host"] = msg["public_host"]
            if msg.get("ext_udp"):
                patch["ext_udp"] = str(msg["ext_udp"])
            if patch:
                write_node_config(patch)
                log("直连地址更新:", patch)
            keys = {k: v for k, v in (msg.get("keys") or {}).items() if v}
            if keys:
                try:
                    st, r = await self.local_json("POST", "/tune", keys)
                    log("key 已下发:", sorted((r or {}).get("applied", {}).keys()) if isinstance(r, dict) else st)
                except Exception as e:
                    log("下发 key 失败:", e)
        elif t == "session_start":
            await self.session_start(msg)
        elif t == "session_end":
            await self.session_end(msg)
        elif t == "shutdown":
            log("门户要求关机:", msg.get("reason"))
            if self.manage and self.server_proc is not None:
                self.server_proc.terminate()
            await asyncio.sleep(0.5)
            os._exit(0)
        elif t == "ping":
            await self.send({"type": "pong"})

    async def send(self, msg):
        if self.ws is not None and not self.ws.closed:
            await self.ws.send_str(json.dumps(msg, ensure_ascii=False))

    async def heartbeat_loop(self):
        while self.ws is not None and not self.ws.closed:
            try:
                st, h = await self.local_json("GET", "/health", timeout=8)
                _, _, util, used = gpu_info()
                await self.send({"type": "heartbeat", "health": h if st == 200 else {"error": "HTTP %s" % st},
                                 "gpu_util": util, "gpu_mem_used": used, "ts": time.time()})
            except Exception as e:
                try:
                    await self.send({"type": "heartbeat", "health": {"error": str(e)}, "ts": time.time()})
                except Exception:
                    pass
            await asyncio.sleep(10)

    async def enroll(self):
        """没有令牌时向门户自动注册。成功 → 令牌写回 agent_config.json；被拉黑（403）不退出：每 5 分钟再试并把原因写进日志，
        管理员解除拉黑后自动接上。其它失败按 5→60s 退避重试。"""
        url = self.portal + "/api/node/enroll"
        wait = 5
        while True:
            name, vram, _, _ = gpu_info()
            body = {"fingerprint": machine_fingerprint(), "instance": instance_key(self.local_port),
                    "hostname": socket.gethostname(), "gpu": name, "vram_mb": vram,
                    "version": AGENT_VERSION, "os": platform.platform(), "name": CFG.get("NODE_NAME", "")}
            try:
                async with self.sess_http.post(url, json=body, timeout=aiohttp.ClientTimeout(total=30)) as r:
                    try:
                        j = await r.json(content_type=None)
                    except Exception:
                        j = {}
                    if r.status == 200 and j.get("token"):
                        self.token = j["token"]
                        save_agent_token(self.token)
                        log("✅ 已自动注册到门户 %s：节点「%s」(#%s)%s" % (self.portal, j.get("name"), j.get("node_id"),
                            "" if j.get("enabled", True) else "（后台已停用，等管理员启用后自动接上）"))
                        return
                    if r.status == 403:
                        wait = 300
                        log("⛔ 门户拒绝注册：这台机器已被管理员拉黑。解除拉黑后会自动接上，%ds 后再试" % wait)
                    else:
                        log("自动注册失败 HTTP %d: %s，%ds 后重试" % (r.status, str(j)[:160], wait))
            except Exception as e:
                log("自动注册失败: %s: %s，%ds 后重试" % (type(e).__name__, str(e)[:160], wait))
            await asyncio.sleep(wait)
            if wait < 300:
                wait = min(wait * 2, 60)

    async def run(self):
        if not self.portal:
            log("❌ 未配置 PORTAL_URL（环境变量或 agent_config.json），无法注册")
            sys.exit(2)
        self.sess_http = aiohttp.ClientSession()
        if not self.token:
            log("没有节点令牌，向门户自动注册…")
            await self.enroll()
        if self.manage:
            self.start_server()
        log("等待本机 server.py 就绪 %s …" % self.local)
        h = await self.wait_local_health()
        if h is None:
            log("❌ server.py 20 分钟未就绪，退出让外层重启")
            sys.exit(3)
        log("server.py 就绪 mode=%s" % h.get("mode"))
        # 断电/异常退出留下的工作区先清掉
        st = _labels(WS_STATE)
        if st.get("sid"):
            await self.session_end({"sid": st["sid"], "user_id": st.get("user_id"), "reason": "agent restart"})
        backoff = 3
        wsurl = self.portal.replace("https://", "wss://").replace("http://", "ws://") + "/api/node/ws"
        while True:
            try:
                async with self.sess_http.ws_connect(wsurl, headers={"Authorization": "Bearer " + self.token},
                                                     heartbeat=25, max_msg_size=64 * 1024 * 1024,
                                                     timeout=aiohttp.ClientWSTimeout(ws_close=10)) as ws:
                    self.ws = ws
                    backoff = 3
                    name, vram, _, _ = gpu_info()
                    cfg = read_node_config()
                    try:
                        local_port = int(self.local.rsplit(":", 1)[1])
                    except Exception:
                        local_port = 0
                    await self.send({"type": "register", "hostname": socket.gethostname(), "gpu": name, "vram_mb": vram,
                                     "fingerprint": machine_fingerprint(), "instance": instance_key(local_port),
                                     "lan_ips": lan_ips(), "local_port": local_port,
                                     "version": AGENT_VERSION, "backend": "vast" if os.environ.get("VAST_CONTAINERLABEL") or os.environ.get("VAST_TCP_PORT_8080") else "installed",
                                     "os": platform.platform(), "http_token": self.http_token,
                                     "public_host": os.environ.get("PUBLIC_IP") or cfg.get("public_host"),
                                     "ext_udp": os.environ.get("VAST_UDP_PORT_8081") or cfg.get("ext_udp"),
                                     "session_id": (self.session or {}).get("sid"), "caps": ["tunnel", "workspace"]})
                    log("已连接门户 %s，注册为 %s / %s" % (self.portal, socket.gethostname(), name))
                    hb = asyncio.ensure_future(self.heartbeat_loop())
                    try:
                        async for m in ws:
                            if m.type == aiohttp.WSMsgType.TEXT:
                                try:
                                    msg = json.loads(m.data)
                                except Exception:
                                    continue
                                asyncio.ensure_future(self.handle(msg))
                            elif m.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
                                break
                    finally:
                        hb.cancel()
                        self.ws = None
                log("门户连接断开，%ds 后重连" % backoff)
            except Exception as e:
                log("连接门户失败: %s: %s，%ds 后重试" % (type(e).__name__, str(e)[:160], backoff))
                # 握手 403 = 令牌无效 / 节点被停用 / 被拉黑 / 被删除。自动注册来的令牌 → 重新注册一次：
                # 被删除会得到新节点；同一机器仍在（停用/拉黑）会拿回同一令牌或 403，并在日志里说明原因，不会绕过停用。
                if isinstance(e, aiohttp.WSServerHandshakeError) and e.status == 403 and not self.token_from_env:
                    log("门户拒绝令牌（节点被停用/拉黑/删除），重新自动注册…")
                    self.token = ""
                    await self.enroll()
                    backoff = 30
            await asyncio.sleep(backoff)
            backoff = min(backoff * 2, 30)


if __name__ == "__main__":
    try:
        asyncio.run(Agent().run())
    except KeyboardInterrupt:
        pass
