# -*- coding: utf-8 -*-
"""让中转给被控端传递「当前是哪个手表会话」。

为什么需要：
    被控端只有一个到中转的长连接。手表断开又连上时，被控端必须知道
    "这是新会话"，才能把上一条连接的状态（鉴权、正在镜像…）清掉。

做法：
    1. 一个设备同时只允许 1 个手表会话（新的来了就踢旧的）
       —— 之前是 3 个，但那样被控端分不清哪段字节属于哪个手表
    2. /v1/agent/poll 的响应头带上 X-Relay-Session: <会话令牌>
       令牌一变，被控端就知道要重开一条会话
"""
from __future__ import annotations

import sys
from pathlib import Path

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

# ---------------- 1. Device 记住当前会话 ----------------

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

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

OLD_INIT = '''        self.bytes_in = 0
        self.bytes_out = 0'''

NEW_INIT = '''        self.bytes_in = 0
        self.bytes_out = 0
        # 当前挂在这台设备上的手表会话令牌。变了就说明手表换了一台/重连了，
        # 被控端据此重开一条会话，避免把上一条的状态带过来。
        self.current_session = ""'''

# ---------------- 2. attach：一个设备只留一个会话 ----------------

OLD_ATTACH = '''            # 一个设备最多留几个会话，超了踢最旧的
            mine = [s for s in self.sessions.values()
                    if s.device_id == device_id]
            if len(mine) >= MAX_SESSIONS_PER_DEVICE:
                mine.sort(key=lambda s: s.last_seen)
                for stale in mine[:len(mine) - MAX_SESSIONS_PER_DEVICE + 1]:
                    self.sessions.pop(stale.token, None)
                    log(f"踢掉旧会话 {stale.token[:8]}…")

            session = Session(device_id, code, name, ip)
            self.sessions[session.token] = session'''

NEW_ATTACH = '''            # 一台设备同时只留一个手表会话。
            # （之前留 3 个，但被控端只有一条到中转的流，分不清哪段字节
            #   属于哪个手表 —— 所以这里改成新的来了就踢旧的。）
            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]}…（同一设备只允许一个）")

            session = Session(device_id, code, name, ip)
            self.sessions[session.token] = session
            device.current_session = session.token
            # 换会话了：把上一条留在队列里的半截数据清掉，
            # 不然新手表会收到上一个手表没取走的残留字节。
            device.to_client.clear()'''

# ---------------- 3. poll 响应头带上会话令牌 ----------------

OLD_POLL = '''        if path == "/v1/agent/poll":
            device = RELAY.device_by_token(self.headers.get("X-Token", ""))
            if device is None:
                self._json(401, {"ok": False, "error": "令牌无效"})
                return
            data = RELAY.agent_take(device, POLL_SECONDS)
            self._send(200, data, "application/octet-stream")
            return'''

NEW_POLL = '''        if path == "/v1/agent/poll":
            device = RELAY.device_by_token(self.headers.get("X-Token", ""))
            if device is None:
                self._json(401, {"ok": False, "error": "令牌无效"})
                return
            session_at_start = device.current_session
            data = RELAY.agent_take(device, POLL_SECONDS)
            self._send(200, data, "application/octet-stream",
                       extra={"X-Relay-Session": device.current_session or "-",
                              "X-Relay-Session-AtStart": session_at_start or "-"})
            return'''

# ---------------- 4. _send 支持额外响应头 ----------------

OLD_SEND = '''    def _send(self, code: int, body: bytes = b"",
              content_type: str = "application/json") -> None:
        self.send_response(code)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()'''

NEW_SEND = '''    def _send(self, code: int, body: bytes = b"",
              content_type: str = "application/json",
              extra: dict | None = None) -> None:
        self.send_response(code)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        for name, value in (extra or {}).items():
            self.send_header(name, value)
        self.end_headers()'''

# ---------------- 5. detached：清掉 current_session ----------------

OLD_BYE = '''        if path == "/v1/agent/bye":
            device = RELAY.device_by_token(self.headers.get("X-Token", ""))
            if device is not None:
                device.last_seen = 0
                log(f"被控端下线：{device.name}")'''

NEW_BYE = '''        if path == "/v1/agent/bye":
            device = RELAY.device_by_token(self.headers.get("X-Token", ""))
            if device is not None:
                device.last_seen = 0
                device.current_session = ""
                log(f"被控端下线：{device.name}")'''


def main() -> int:
    text = RELAY.read_text(encoding="utf-8")
    ok = True
    for old, new, tag in ((OLD_SLOTS, NEW_SLOTS, "Device 加 current_session 槽位"),
                          (OLD_INIT, NEW_INIT, "初始化 current_session"),
                          (OLD_ATTACH, NEW_ATTACH, "一台设备只留一个会话"),
                          (OLD_POLL, NEW_POLL, "poll 带会话令牌"),
                          (OLD_SEND, NEW_SEND, "_send 支持额外响应头"),
                          (OLD_BYE, NEW_BYE, "下线时清会话")):
        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())