# -*- coding: utf-8 -*-
"""第一步：让手表知道"电脑在哪个局域网、开在哪个端口"。

为什么需要：
    用户要求 —— 手表连上之后，如果发现电脑和自己**在同一个网络**，
    就切换成局域网直连（快得多）；不在同一网络就继续用中转。

    要判断"同一个网络"，手表得知道电脑的局域网 IP。
    所以被控端登记时要把自己的局域网 IP 和端口报给中转，
    中转再转给手表。

改动：
    1. windows/lanremote/relay.py  登记时带上 lanIps 和 port
    2. relay/relay.py              Device 存下来，并在设备信息里返回
"""
from __future__ import annotations

import sys
from pathlib import Path

ROOT = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote")
CLIENT = ROOT / "windows" / "lanremote" / "relay.py"
SERVER = ROOT / "relay" / "relay.py"

# ---------------- 1. 被控端：算自己的局域网 IP ----------------

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

CLIENT_NEW = '''    def _lan_ips(self) -> list:
        """本机在局域网里的 IP（可能有多张网卡）。

        手表靠这个判断"电脑是不是和我同一个网络"。
        拿不到就返回空列表 —— 中转那边会当成"不知道"，
        手表就继续走中转，不会出错。
        """
        found = []
        try:
            probe = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            try:
                # 不会真的发包，只是让系统选出默认出口网卡
                probe.connect(("8.8.8.8", 80))
                addr = probe.getsockname()[0]
                if addr and not addr.startswith("127."):
                    found.append(addr)
            finally:
                probe.close()
        except OSError:
            pass
        try:
            for info in socket.getaddrinfo(socket.gethostname(), None,
                                           socket.AF_INET):
                addr = info[4][0]
                if (addr and not addr.startswith("127.")
                        and not addr.startswith("169.254.")
                        and addr not in found):
                    found.append(addr)
        except OSError:
            pass
        return found

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

CLIENT_PAYLOAD_OLD = '''        payload = json.dumps({
            "deviceId": self.device_id,
            "code": self.code,
            "name": self.device_name or self.device_id,
            "version": _version(),
        }, ensure_ascii=False).encode("utf-8")'''

CLIENT_PAYLOAD_NEW = '''        try:
            listen_port = int(self.store.get("port") or 0)
        except (TypeError, ValueError):
            listen_port = 0
        payload = json.dumps({
            "deviceId": self.device_id,
            "code": self.code,
            "name": self.device_name or self.device_id,
            "version": _version(),
            # 手表用这两个判断要不要切成局域网直连
            "lanIps": self._lan_ips(),
            "port": listen_port,
        }, ensure_ascii=False).encode("utf-8")'''

# ---------------- 2. 中转：Device 存字段 ----------------

SRV_SLOTS_OLD = '''    __slots__ = ("device_id", "code", "name", "version", "token",
                 "last_seen", "connected_at", "to_client", "to_agent",
                 "event", "bytes_in", "bytes_out", "current_session",
                 "active_session")'''

SRV_SLOTS_NEW = '''    __slots__ = ("device_id", "code", "name", "version", "token",
                 "last_seen", "connected_at", "to_client", "to_agent",
                 "event", "bytes_in", "bytes_out", "current_session",
                 "active_session", "lan_ips", "port")'''

SRV_INIT_OLD = '''        self.token = secrets.token_hex(24)
        self.last_seen = time.time()'''

SRV_INIT_NEW = '''        self.token = secrets.token_hex(24)
        # 这台电脑的局域网 IP 和监听端口。
        # 手表拿它判断"要不要切成局域网直连"。
        self.lan_ips: list = []
        self.port: int = 0
        self.last_seen = time.time()'''

# 注册时接收
SRV_REG_OLD = '''        name = str(body.get("name") or "").strip()[:40] or "未命名设备"
        version = str(body.get("version") or "").strip()[:20]'''

SRV_REG_NEW = '''        name = str(body.get("name") or "").strip()[:40] or "未命名设备"
        version = str(body.get("version") or "").strip()[:20]
        # 这台电脑的局域网 IP + 端口（手表判断要不要走局域网直连）
        raw_ips = body.get("lanIps")
        lan_ips = []
        if isinstance(raw_ips, list):
            for item in raw_ips[:8]:
                text = str(item).strip()
                # 只收看起来像 IPv4 的，别把乱七八糟的东西转给手表
                parts = text.split(".")
                if len(parts) == 4 and all(
                        p.isdigit() and 0 <= int(p) <= 255 for p in parts):
                    lan_ips.append(text)
        try:
            listen_port = int(body.get("port") or 0)
        except (TypeError, ValueError):
            listen_port = 0
        if not (0 < listen_port < 65536):
            listen_port = 0'''

# 存到 Device 上（在设备对象拿到之后）
SRV_STORE_ANCHOR = '''            device.last_seen = time.time()'''

SRV_STORE_NEW = '''            device.last_seen = time.time()
            # 每次登记都刷新局域网信息（IP 可能变，端口也可能改）
            device.lan_ips = lan_ips
            device.port = listen_port'''

# 快照里带给手表
SRV_SNAP_OLD = '''                "code": device.code,
                "name": device.name,'''

SRV_SNAP_NEW = '''                "code": device.code,
                "name": device.name,
                # 手表用这两个判断能不能走局域网直连
                "lanIps": list(device.lan_ips),
                "port": device.port,'''


def patch(path: Path, pairs, label: str) -> bool:
    text = path.read_text(encoding="utf-8")
    ok = True
    for old, new, tag in pairs:
        if old not in text:
            print(f"  [{label}] [跳过] {tag} —— 锚点没找到")
            ok = False
            continue
        text = text.replace(old, new, 1)
        print(f"  [{label}] {tag}")
    path.write_text(text, encoding="utf-8", newline="\n")
    return ok


def main() -> int:
    ok = True

    ok &= patch(CLIENT, [
        (CLIENT_ANCHOR, CLIENT_NEW, "_lan_ips() 方法"),
        (CLIENT_PAYLOAD_OLD, CLIENT_PAYLOAD_NEW, "登记带上 lanIps/port"),
    ], "被控端")

    ok &= patch(SERVER, [
        (SRV_SLOTS_OLD, SRV_SLOTS_NEW, "Device 加字段"),
        (SRV_INIT_OLD, SRV_INIT_NEW, "Device.__init__"),
        (SRV_REG_OLD, SRV_REG_NEW, "注册接收 lanIps/port"),
        (SRV_STORE_ANCHOR, SRV_STORE_NEW, "存到 Device"),
        (SRV_SNAP_OLD, SRV_SNAP_NEW, "快照里返回"),
    ], "中转")

    import ast
    for label, path in (("被控端", CLIENT), ("中转", SERVER)):
        try:
            ast.parse(path.read_text(encoding="utf-8"))
            print(f"  {label} 语法 OK")
        except SyntaxError as exc:
            print(f"  {label} 语法错误: 行 {exc.lineno}: {exc.msg}")
            ok = False
    return 0 if ok else 1


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