# -*- coding: utf-8 -*-
r'''中转管理页新增两个功能（中转端逻辑）。

== 功能一：允许/禁止新设备接入 ==

一个开关，关掉之后**只有已经登记过的设备能继续上线**，
没见过的新设备一律拒绝。

用途：不想让陌生人拿配对码乱试、或者临时把中转"关起来"。

实现要点：
  * 标志位存在 relay_state.json 里，重启不丢
  * 判断"新设备"以 deviceId 为准（它装机时生成、重装也不变），
    不是看配对码 —— 老设备换码也必须能进来
  * 被拒绝时回一个明确的标记 blocked=True，
    让客户端能区分"被管理员挡了"和"配对码冲突"，
    不要用中文错误文案去匹配（文案一改就坏）

== 功能二：修改任意设备的配对码 ==

管理页上每台设备一个"改配对码"按钮，填新码就生效。

实现要点：
  * 新码必须是 6 位数字
  * 新码不能已经被别的设备占用
  * 改的时候要同时更新 codes 映射和 device.code，
    并清掉该设备**旧码**的映射 —— 否则旧码会变成孤儿条目，
    仪表盘上"配对码占用"会虚高
  * 改完立刻 _save()，重启后仍然是新码
  * 该设备当前的手表会话要踢掉（会话是绑在旧码上的），
    否则手表还拿着旧会话以为连着，实际电脑已经换了码
'''
from __future__ import annotations

import sys
from pathlib import Path

SRC = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote\relay\relay.py")


def patch(text: str, old: str, new: str, label: str) -> tuple:
    if old not in text:
        print("  ✗ 没找到：%s" % label)
        return text, False
    if text.count(old) != 1:
        print("  ! 出现 %d 次：%s" % (text.count(old), label))
        return text, False
    print("  ✓ %s" % label)
    return text.replace(old, new, 1), True


def main() -> int:
    text = SRC.read_text(encoding="utf-8")
    before = text.count("\n") + 1
    ok_all = True

    # ---------------------------------------------------------- 1. 字段
    old = """        self.retired: dict = {}
        self.started = time.time()"""
    new = """        self.retired: dict = {}
        # 是否允许新设备接入。False 时只有已登记过的 deviceId 能上线。
        self.allow_new = True
        self.started = time.time()"""
    text, ok = patch(text, old, new, "加 allow_new 字段")
    ok_all = ok_all and ok

    # ---------------------------------------------------------- 2. _load
    old = """        for item in raw.get("codes", []):
            device_id = str(item.get("deviceId", ""))
            code = str(item.get("code", ""))
            if device_id and CODE_RE.match(code):
                self.codes[code] = device_id
                log(f"恢复配对码绑定：设备 {device_id[:12]}… -> {code}")"""
    new = """        # 是否允许新设备接入（重启不丢）
        self.allow_new = bool(raw.get("allowNew", True))
        for item in raw.get("codes", []):
            device_id = str(item.get("deviceId", ""))
            code = str(item.get("code", ""))
            if device_id and CODE_RE.match(code):
                self.codes[code] = device_id
                log(f"恢复配对码绑定：设备 {device_id[:12]}… -> {code}")"""
    text, ok = patch(text, old, new, "_load 恢复 allow_new")
    ok_all = ok_all and ok

    # ---------------------------------------------------------- 3. _save
    old = """            payload = {"codes": [
                {"deviceId": did, "code": code}
                for code, did in sorted(self.codes.items())
            ]}"""
    new = """            payload = {
                "allowNew": bool(self.allow_new),
                "codes": [
                    {"deviceId": did, "code": code}
                    for code, did in sorted(self.codes.items())
                ],
            }"""
    text, ok = patch(text, old, new, "_save 持久化 allow_new")
    ok_all = ok_all and ok

    # ---------------------------------------------------------- 4. 拒绝新设备
    old = """        with self.lock:
            known = self.codes.get(code)
            if known and known != device_id:"""
    new = """        with self.lock:
            # 管理员关掉了"允许新设备接入"时，
            # 只放行**已经登记过**的设备。
            #
            # 注意判断依据是 deviceId（装机时生成、重装也不变），
            # **不是配对码** —— 否则管理员给老设备改了码，
            # 那台设备就被自己的开关挡在外面了。
            known_device = device_id in self.devices
            if not known_device and not self.allow_new:
                log(f"新设备接入已被关闭，拒绝 {name} ({device_id[:12]}…)")
                return {"ok": False,
                        "blocked": True,
                        "error": "中转站当前不接受新设备接入"}

            known = self.codes.get(code)
            if known and known != device_id:"""
    text, ok = patch(text, old, new, "register 拒绝新设备")
    ok_all = ok_all and ok

    # ---------------------------------------------------------- 5. 管理方法
    old = """    def device_by_token(self, token: str) -> Device | None:"""
    new = """    # ---------------- 管理操作 ----------------

    def set_allow_new(self, allow: bool) -> dict:
        \"\"\"开关"允许新设备接入"。\"\"\"
        with self.lock:
            self.allow_new = bool(allow)
            self._save()
            log(f"管理员把「允许新设备接入」设为 {self.allow_new}")
            return {"ok": True, "allowNew": self.allow_new}

    def set_device_code(self, device_id: str, new_code: str) -> dict:
        \"\"\"给一台设备换配对码。

        要点：
          * 新码必须是 6 位数字
          * 不能已经被**别的**设备占用
          * 旧码的映射必须删掉，否则会变成孤儿条目，
            管理页上"配对码占用"的数字会虚高
          * 该设备挂着的会话要踢掉 —— 会话是绑在码上的，
            不踢的话手表还拿着旧会话以为连着
        \"\"\"
        device_id = str(device_id or "").strip()
        new_code = str(new_code or "").strip()
        if not CODE_RE.match(new_code):
            return {"ok": False, "error": "配对码必须是 6 位数字"}
        with self.lock:
            device = self.devices.get(device_id)
            if device is None:
                return {"ok": False, "error": "没有这台设备（可能已下线）"}

            owner = self.codes.get(new_code)
            if owner and owner != device_id:
                return {"ok": False,
                        "error": "这个配对码已经被别的设备占用了"}

            old_code = device.code
            if old_code == new_code:
                return {"ok": True, "code": new_code, "unchanged": True}

            # 删掉旧码映射（连同可能存在的别名）
            for code, owner_id in list(self.codes.items()):
                if owner_id == device_id:
                    self.codes.pop(code, None)
            self.codes[new_code] = device_id
            device.code = new_code

            # 会话绑在码上，旧会话作废，让手表重新配对
            kicked = 0
            for sid in [s.session_id for s in self.sessions.values()
                        if s.device_id == device_id]:
                self.sessions.pop(sid, None)
                kicked += 1

            self._save()
            log(f"管理员把设备 {device.name} 的配对码 "
                f"{old_code} 改成 {new_code}（踢掉 {kicked} 个会话）")
            return {"ok": True, "code": new_code,
                    "oldCode": old_code, "kicked": kicked}

    def device_by_token(self, token: str) -> Device | None:"""
    text, ok = patch(text, old, new, "加 set_allow_new / set_device_code")
    ok_all = ok_all and ok

    if not ok_all:
        print()
        print("  有锚点没匹配，没有写回")
        return 1

    SRC.write_text(text, encoding="utf-8", newline="\n")
    print()
    print("  已写回：%d -> %d 行" % (before, text.count("\n") + 1))
    return 0


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