# -*- coding: utf-8 -*-
"""中转：区分「会话过期」和「被新连接替换」。

问题（联调时暴露的）：
    一台设备只允许一个手表会话。新手表接入 -> 踢掉旧的。
    但旧的被踢后拿到 401，客户端会**自动重连** -> 又把新的踢掉
    -> 互相踢，永远握不上手，表现为两端一直超时。

修法：
    被踢掉的会话标记成「已替换」。它再来问就返回 409 + replaced=true，
    客户端看到这个就**老老实实断开**（新手表赢了），不再抢。
    只有中转重启导致的令牌失效（401）才触发自动重连。

用户看到的效果：
    新表连上 -> 旧表自动退出并提示"已被另一台设备接管"，不会来回抢。
"""
from __future__ import annotations

import sys
from pathlib import Path

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

# ---------------- 1. Relay 加 retired 记录 ----------------

OLD_INIT = '''        self.attach_log: dict = {}         # ip -> [时间戳]'''

NEW_INIT = '''        self.attach_log: dict = {}         # ip -> [时间戳]
        # 被新连接顶替掉的会话：token -> 原因。
        # 留着是为了能明确告诉旧客户端"你是被顶掉的"，
        # 而不是含糊的 401 —— 否则它会以为自己只是过期了，然后抢回来。
        self.retired: dict = {}'''

# ---------------- 2. 踢的时候记下来 ----------------

OLD_KICK = '''            for stale_token, stale in list(self.sessions.items()):
                if stale.device_id == device_id:
                    self.sessions.pop(stale_token, None)
                    log(f"踢掉旧会话 {stale_token[:8]}…（同一设备只允许一个）")'''

NEW_KICK = '''            for stale_token, stale in list(self.sessions.items()):
                if stale.device_id == device_id:
                    self.sessions.pop(stale_token, None)
                    self.retired[stale_token] = "replaced"
                    log(f"踢掉旧会话 {stale_token[:8]}…（同一设备只允许一个）")
            # 清掉太老的顶替记录，别让字典涨
            if len(self.retired) > 200:
                for old_token in list(self.retired)[:100]:
                    self.retired.pop(old_token, None)'''

# ---------------- 3. 查会话时区分 ----------------

OLD_LOOKUP = '''    def session_by_token(self, token: str) -> Session | None:
        if not token:
            return None
        with self.lock:
            session = self.sessions.get(token)
            if session is None:
                return None
            session.last_seen = time.time()
            return session'''

NEW_LOOKUP = '''    def session_by_token(self, token: str) -> Session | None:
        if not token:
            return None
        with self.lock:
            session = self.sessions.get(token)
            if session is None:
                return None
            session.last_seen = time.time()
            return session

    def session_replaced(self, token: str) -> bool:
        """这个令牌是不是「被新连接顶掉」的（而不是过期/不存在）。"""
        if not token:
            return False
        with self.lock:
            return self.retired.get(token) == "replaced"'''

# ---------------- 4. 接口返回 409 ----------------

OLD_POLL = '''        if path == "/v1/client/poll":
            session = RELAY.session_by_token(self.headers.get("X-Session", ""))
            if session is None:
                self._json(401, {"ok": False, "error": "会话无效"})
                return'''

NEW_POLL = '''        if path == "/v1/client/poll":
            token = self.headers.get("X-Session", "")
            session = RELAY.session_by_token(token)
            if session is None:
                if RELAY.session_replaced(token):
                    # 明确告诉它"你是被顶掉的"，客户端就不会抢回来了
                    self._json(409, {"ok": False, "replaced": True,
                                     "error": "已被另一台手表接管"})
                    return
                self._json(401, {"ok": False, "error": "会话无效"})
                return'''

OLD_SEND = '''        if path == "/v1/client/send":
            session = RELAY.session_by_token(self.headers.get("X-Session", ""))
            if session is None:
                self._json(401, {"ok": False, "error": "会话无效"})
                return'''

NEW_SEND = '''        if path == "/v1/client/send":
            token = self.headers.get("X-Session", "")
            session = RELAY.session_by_token(token)
            if session is None:
                if RELAY.session_replaced(token):
                    self._json(409, {"ok": False, "replaced": True,
                                     "error": "已被另一台手表接管"})
                    return
                self._json(401, {"ok": False, "error": "会话无效"})
                return'''

# ---------------- 5. detach 时也清掉记录 ----------------

OLD_DETACH = '''            if session is not None:
                relay_sessions = getattr(RELAY, "sessions", {})
                relay_sessions.pop(session.token, None)'''

NEW_DETACH = '''            if session is not None:
                relay_sessions = getattr(RELAY, "sessions", {})
                relay_sessions.pop(session.token, None)
                RELAY.retired.pop(session.token, None)'''


def main() -> int:
    text = RELAY.read_text(encoding="utf-8")
    ok = True
    for old, new, tag in ((OLD_INIT, NEW_INIT, "Relay 加 retired"),
                          (OLD_KICK, NEW_KICK, "踢的时候记录原因"),
                          (OLD_LOOKUP, NEW_LOOKUP, "session_replaced 查询"),
                          (OLD_POLL, NEW_POLL, "poll 返回 409"),
                          (OLD_SEND, NEW_SEND, "send 返回 409"),
                          (OLD_DETACH, NEW_DETACH, "detach 清理")):
        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")

    import ast
    try:
        ast.parse(text)
        print("  语法 OK")
    except SyntaxError as exc:
        print(f"  语法错误: {exc}")
        ok = False
    return 0 if ok else 1


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