# -*- coding: utf-8 -*-
r'''把 swap_admin_page.py 误删的 4 个顶层定义补回来。

丢的（git 里还在）：
    _esc      2006    HTML 转义
    _kb       2011    字节数格式化
    _uptime   2019    运行时长格式化
    Server    2034    ThreadingHTTPServer 子类 —— main() 要用它，
                      少了它就是 NameError，relay 起不来

做法：从 git 取下这一整段，去掉里面已经存在的部分（比如 main），
      插到当前文件的 main() 前面。
'''
from __future__ import annotations

import ast
import re
import subprocess
import sys
from pathlib import Path

ROOT = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote")
RELAY = ROOT / "relay" / "relay.py"
GIT = r"C:\Program Files\Git\bin\git.exe"


def top_defs(source: str) -> set:
    tree = ast.parse(source)
    names = set()
    for node in tree.body:
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef,
                             ast.ClassDef)):
            names.add(node.name)
        elif isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name):
                    names.add(target.id)
    return names


def main() -> int:
    old = subprocess.run([GIT, "show", "HEAD:relay/relay.py"],
                         cwd=str(ROOT), capture_output=True, text=True,
                         encoding="utf-8", errors="replace").stdout
    cur = RELAY.read_text(encoding="utf-8")

    have = top_defs(cur)
    missing = {"_esc", "_kb", "_uptime", "Server"} - have
    if not missing:
        print("  没有缺的了")
        return 0
    print("  缺: %s" % ", ".join(sorted(missing)))

    # 从 git 里按段摘出来
    lines = old.splitlines()
    blocks = []
    index = 0
    while index < len(lines):
        line = lines[index]
        match = re.match(r"^(def|class)\s+(\w+)", line)
        if match and match.group(2) in missing:
            # 收集这一段（直到下一个顶层定义或文件尾）
            block = [line]
            index += 1
            while index < len(lines):
                nxt = lines[index]
                if re.match(r"^(def|class)\s+\w+", nxt) or nxt.startswith("if __name__"):
                    break
                block.append(nxt)
                index += 1
            # 去掉尾部多余空行
            while block and not block[-1].strip():
                block.pop()
            blocks.append("\n".join(block))
            print("    摘到 %s（%d 行）" % (match.group(2), len(block)))
            continue
        index += 1

    if not blocks:
        print("  [失败] 一段都没摘到")
        return 1

    # 插到 main() 前面
    marker = "\ndef main("
    position = cur.find(marker)
    if position < 0:
        print("  [失败] 找不到 main()")
        return 1

    header = ("\n\n# ======================================================================\n"
              "# 下面这几个是页面/启动辅助（重做管理页时被误删，从 git 恢复）\n"
              "# ======================================================================\n\n")
    text = (cur[:position + 1] + header + "\n\n".join(blocks) + "\n"
            + cur[position:])

    RELAY.write_text(text, encoding="utf-8", newline="\n")

    try:
        ast.parse(text)
        print("  relay.py 语法 OK")
    except SyntaxError as exc:
        print("  [失败] 行 %s: %s" % (exc.lineno, exc.msg))
        return 1

    now = top_defs(text)
    still = missing - now
    if still:
        print("  [失败] 还缺: %s" % ", ".join(sorted(still)))
        return 1
    print("  ✓ 4 个都补回来了")

    print()
    print("  复核（关键定义是否都在）:")
    for name in ("_esc", "_kb", "_uptime", "Server", "Handler", "main",
                 "admin_key", "public_url", "_list_downloads"):
        print("    %-18s %s" % (name, "有" if name in now else "**缺**"))
    return 0


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