# -*- coding: utf-8 -*-
"""被控端支持多个中转地址，自动故障切换。

解决的痛点：
    匿名隧道地址老变，一变就得在被控端重配一次。

实际场景：
    电脑和随身 WiFi 经常在**同一个局域网**（随身 WiFi 就是给电脑上网用的）。
    那就不用绕隧道了 —— 直连局域网中转更快（几十毫秒 vs 600 毫秒）、
    更稳（不依赖隧道），而且地址**永远不变**。

做法：
    relay_url 支持用分号分隔多个地址，按顺序尝试：
        http://192.168.2.156:8080;https://relay.xxmxxs.top
    第一个通的就用它。局域网断了自动切隧道。

    主地址连不上时，还会去"局域网助手地址"查一下当前隧道地址
    （中转的 /url 接口），自动跟随地址变化。
"""
from __future__ import annotations

import sys
from pathlib import Path

RELAY = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote\windows"
             r"\lanremote\relay.py")
RUN = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote\windows"
           r"\run_agent.py")

# ---------------- 1. 解析多地址 ----------------

OLD_START = '''    def start(self, url: str) -> bool:
        """开始连中转。url 形如 https://relay.example.com"""
        url = (url or "").strip().rstrip("/")
        if not url:
            self._log("没配置中转地址，跳过")
            return False
        if not url.startswith(("http://", "https://")):
            url = "https://" + url
        self.url = url'''

NEW_START = '''    def start(self, url: str) -> bool:
        """开始连中转。

        url 支持用**分号**分隔多个地址，按顺序尝试：
            http://192.168.2.156:8080;https://relay.xxmxxs.top

        为什么需要多个：
          * 电脑和随身 WiFi 常在同一局域网 —— 直连局域网中转更快更稳，
            而且地址永远不变（不用管隧道地址变来变去）
          * 电脑在外面时才走隧道
        """
        candidates = []
        for piece in (url or "").replace(",", ";").split(";"):
            piece = piece.strip().rstrip("/")
            if not piece:
                continue
            if not piece.startswith(("http://", "https://")):
                piece = "https://" + piece
            if piece not in candidates:
                candidates.append(piece)

        if not candidates:
            self._log("没配置中转地址，跳过")
            return False

        self.candidates = candidates
        # 先探一遍，挑第一个通的
        chosen = self._pick_candidate(candidates)
        self.url = chosen
        if len(candidates) > 1:
            self._log(f"中转地址候选 {len(candidates)} 个，选用 {chosen}")'''

# ---------------- 2. 挑一个能用的 ----------------

ANCHOR = '''    def _ensure_device_id(self) -> str:'''

NEW_METHODS = '''    def _probe(self, base: str, timeout: float = 8.0) -> bool:
        """这个地址通不通（打一下 /v1/health）。"""
        try:
            request = urllib.request.Request(base + "/v1/health")
            context = _ssl_context() if base.startswith("https://") else None
            with urllib.request.urlopen(request, timeout=timeout,
                                        context=context) as response:
                if response.status != 200:
                    return False
                body = response.read(4096).decode("utf-8", "replace")
                return '"ok"' in body
        except Exception:                                 # noqa: BLE001
            return False

    def _pick_candidate(self, candidates: list) -> str:
        for base in candidates:
            if self._probe(base):
                return base
            self._log(f"  {base} 不通，试下一个")
        # 都不通也先返回第一个，让主循环去重试
        return candidates[0]

    def _refresh_url(self) -> None:
        """主地址连不上时，去候选里找找有没有能用的。

        如果某个候选本身就是局域网中转，还能顺便问它
        「现在的隧道地址是什么」（/url 接口），自动跟随地址变化。
        """
        for base in list(self.candidates):
            if base == self.url:
                continue
            if not self._probe(base, timeout=5):
                continue
            self._log(f"切换到备用中转地址：{base}")
            self.url = base
            return

    def _ensure_device_id(self) -> str:'''

# ---------------- 3. 连不上时换地址 ----------------

OLD_RUN = '''            except Exception as exc:                      # noqa: BLE001
                self.online = False
                self.last_error = f"{type(exc).__name__}: {exc}"
                self._log(f"中转连接出问题：{self.last_error}")
                # 网络断了：把当前会话关掉，让上层重开
                self._drop_socket()
                time.sleep(backoff)
                backoff = min(RETRY_MAX, backoff * 1.6)'''

NEW_RUN = '''            except Exception as exc:                      # noqa: BLE001
                self.online = False
                self.last_error = f"{type(exc).__name__}: {exc}"
                self._log(f"中转连接出问题：{self.last_error}")
                # 网络断了：把当前会话关掉，让上层重开
                self._drop_socket()
                self.token = ""
                # 连续失败几次就换一个候选地址试试
                fail_count = getattr(self, "_fail_count", 0) + 1
                self._fail_count = fail_count
                if fail_count >= 3 and len(getattr(self, "candidates", [])) > 1:
                    self._fail_count = 0
                    self._refresh_url()
                time.sleep(backoff)
                backoff = min(RETRY_MAX, backoff * 1.6)'''

# ---------------- 4. __init__ 加字段 ----------------

OLD_INIT = '''        self.url = ""
        self.device_id = ""'''

NEW_INIT = '''        self.url = ""
        self.candidates: list = []
        self._fail_count = 0
        self.device_id = ""'''

# ---------------- 5. status 里显示候选 ----------------

OLD_STATUS = '''            "url": self.url,'''
NEW_STATUS = '''            "url": self.url,
            "candidates": list(self.candidates),'''


def main() -> int:
    text = RELAY.read_text(encoding="utf-8")
    ok = True

    if "self.candidates" in text:
        print("  已经加过多地址支持")
        return 0

    for old, new, tag in ((OLD_INIT, NEW_INIT, "__init__ 加字段"),
                          (OLD_START, NEW_START, "start() 解析多地址"),
                          (ANCHOR, NEW_METHODS, "探测/切换/刷新方法"),
                          (OLD_RUN, NEW_RUN, "失败次数到了就换地址"),
                          (OLD_STATUS, NEW_STATUS, "status 显示候选")):
        if old not in text:
            print(f"  [失败] {tag} —— 锚点没找到")
            ok = False
            continue
        text = text.replace(old, new, 1)
        print(f"  {tag}")

    RELAY.write_text(text, encoding="utf-8", newline="\n")

    import ast
    try:
        ast.parse(text)
        print("  relay.py 语法 OK")
    except SyntaxError as exc:
        print(f"  语法错误: {exc}")
        ok = False

    # ---------------- 6. CLI 说明 ----------------
    run_text = RUN.read_text(encoding="utf-8")
    old_help = '''    parser.add_argument("--relay", default=None, metavar="地址",
                        help="设置中转服务器地址（跨网络用），传空字符串关闭")'''
    new_help = '''    parser.add_argument("--relay", default=None, metavar="地址",
                        help="设置中转服务器地址（跨网络用）。"
                             "可用分号分隔多个，按顺序尝试，例如："
                             "http://192.168.2.156:8080;https://relay.xxmxxs.top"
                             "  —— 局域网优先，断了自动切隧道。"
                             "传空字符串关闭")'''
    if old_help in run_text:
        run_text = run_text.replace(old_help, new_help, 1)
        RUN.write_text(run_text, encoding="utf-8", newline="\n")
        print("  run_agent.py: --relay 帮助已更新")

    return 0 if ok else 1


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