# -*- coding: utf-8 -*-
"""分块累积：加总体积上限 + 顺序校验。

`MAX_LINE_BYTES` 只管单行，累积体积没上限 —— 已配对设备可以一直发 chunk
把内存吃干（最后 bytes() 还要再翻一倍）。而且原来按到达顺序直接拼，
seq 算出来却从不校验，丢块/乱序会静默产出坏文件还回 ok:True。
"""
from __future__ import annotations

import sys
from pathlib import Path

PROTOCOL = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote\windows"
                r"\lanremote\protocol.py")
SERVER = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote\windows"
              r"\lanremote\server.py")

OLD = '''class ChunkAssembler:
    """接收端：按 seq 收集分块，拼回完整数据。"""

    def __init__(self):
        self._buf = bytearray()
        self._expected = 0

    def feed(self, seq: int, is_last: bool, data: bytes) -> bytes | None:
        """返回完整数据（当收到最后一块时），否则 None。"""
        if seq == 0:
            self._buf = bytearray()
        self._buf.extend(data)
        self._expected = seq + 1
        if is_last:
            result = bytes(self._buf)
            self._buf = bytearray()
            return result
        return None'''

NEW = '''class ChunkTooLarge(Exception):
    """分块累积超过上限，或顺序不对。"""


class ChunkAssembler:
    """接收端：按 seq 收集分块，拼回完整数据。

    MAX_LINE_BYTES 只管单行，累积体积必须另外设上限 ——
    否则已配对设备可以一直发 chunk 把内存吃干（最后 bytes() 还要再翻一倍）。
    """

    def __init__(self, limit: int = MAX_TOTAL_BYTES):
        self._buf = bytearray()
        self._expected = 0
        self._limit = int(limit)

    def feed(self, seq: int, is_last: bool, data: bytes) -> bytes | None:
        """返回完整数据（当收到最后一块时），否则 None。

        顺序不对或超限时抛 ChunkTooLarge，不静默拼出坏文件。
        """
        if seq == 0:
            self._buf = bytearray()
            self._expected = 0
        if seq != self._expected:
            self._buf = bytearray()
            self._expected = 0
            raise ChunkTooLarge(
                f"分块顺序不对（收到 {seq}，应为 {self._expected}）")
        if len(self._buf) + len(data) > self._limit:
            self._buf = bytearray()
            self._expected = 0
            raise ChunkTooLarge(f"文件超过上限（{self._limit // 1048576} MB）")
        self._buf.extend(data)
        self._expected = seq + 1
        if is_last:
            result = bytes(self._buf)
            self._buf = bytearray()
            self._expected = 0
            return result
        return None

    def reset(self) -> None:
        self._buf = bytearray()
        self._expected = 0'''

CONST_OLD = "CHUNK_SIZE = 256 * 1024"
CONST_NEW = ("CHUNK_SIZE = 256 * 1024\n"
             "# 单次传输（文件 / 图片 / 音频）的总体积上限\n"
             "MAX_TOTAL_BYTES = 256 * 1024 * 1024")


def main() -> int:
    ok = True

    text = PROTOCOL.read_text(encoding="utf-8")
    if OLD not in text:
        print("  [失败] protocol.py ChunkAssembler 锚点没找到")
        ok = False
    else:
        text = text.replace(OLD, NEW, 1)
        print("  protocol.py / ChunkAssembler: 加体积上限 + 顺序校验")
    if CONST_OLD in text and "MAX_TOTAL_BYTES" not in text:
        text = text.replace(CONST_OLD, CONST_NEW, 1)
        print("  protocol.py / MAX_TOTAL_BYTES 常量: OK")
    if ok:
        PROTOCOL.write_text(text, encoding="utf-8")

    # ---- server.py 要接住 ChunkTooLarge ----
    stext = SERVER.read_text(encoding="utf-8")
    old_feed = "            blob = assembler.feed(seq, last, data)"
    new_feed = ("            try:\n"
                "                blob = assembler.feed(seq, last, data)\n"
                "            except P.ChunkTooLarge as exc:\n"
                "                session.buffers.pop(cmd, None)\n"
                "                raise ActionError(str(exc)) from exc")
    if old_feed in stext:
        stext = stext.replace(old_feed, new_feed, 1)
        SERVER.write_text(stext, encoding="utf-8")
        print("  server.py / 接住 ChunkTooLarge: OK")
    else:
        print("  [跳过] server.py assembler.feed 锚点没找到")

    import ast
    for label, path in (("protocol.py", PROTOCOL), ("server.py", SERVER)):
        try:
            ast.parse(path.read_text(encoding="utf-8"))
            print(f"  {label}: 语法 OK")
        except SyntaxError as exc:
            print(f"  {label}: 语法错误 {exc}")
            ok = False
    return 0 if ok else 1


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