# -*- coding: utf-8 -*-
"""给测试脚本也加上 User-Agent。

背景：被控端代码已经带了 LanRemote/1.4.0（Cloudflare 的
Browser Integrity Check 会拦 Python 默认 UA，返回 403 + error 1010）。
但测试脚本自己也是用 urllib 发请求的，所以经固定域名跑就全 403。

做法：给每个测试脚本在最前面加一个 _UA 常量，
然后把所有 urllib.request.Request(...) 和 urlopen(...) 都带上它。
"""
from __future__ import annotations

import re
import sys
from pathlib import Path

TESTS = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote\_tests")

FILES = [
    "test_relay_e2e.py",
    "test_relay_local.py",
    "test_relay_admin_login.py",
]

UA_BLOCK = '''

# Cloudflare 的 Browser Integrity Check 会拦 Python 默认 UA
# （返回 403 + error code 1010），所以请求都要带上这个。
_UA = {"User-Agent": "LanRemote/1.4.0"}
'''

# 给 Request(...) 加 headers
REQ_PATTERNS = [
    # Request(url, data=body, method="POST")
    (re.compile(r"urllib\.request\.Request\(([^()]*?),\s*(data=[^()]*?),\s*"
                r"(method=\"POST\")\)"),
     r'urllib.request.Request(\1, \2, \3, headers=_UA)'),
    # Request(url)
    (re.compile(r"urllib\.request\.Request\(([^()]*?)\)"),
     r'urllib.request.Request(\1, headers=_UA)'),
]


def patch_file(path: Path) -> int:
    text = path.read_text(encoding="utf-8")
    if "_UA" in text:
        print(f"  {path.name}: 已经加过")
        return 0

    changes = 0

    # 1. 每个 Request( 都加 headers=_UA
    #    先处理带 method= 的（更长的模式），再处理简单的
    def add_headers(match):
        nonlocal changes
        inner = match.group(1)
        if "headers=" in inner:
            return match.group(0)
        changes += 1
        return f"urllib.request.Request({inner}, headers=_UA)"

    # 匹配 Request( ... ) 的一层括号内容
    text = re.sub(
        r"urllib\.request\.Request\(((?:[^()]|\([^()]*\))*?)\)(?!\s*\))",
        add_headers, text)

    # 2. 裸 urlopen(URL + path, ...) 补 headers
    def add_urlopen_headers(match):
        nonlocal changes
        inner = match.group(1)
        if "headers=" in inner or "_UA" in inner:
            return match.group(0)
        changes += 1
        return f"urllib.request.urlopen({inner}, headers=_UA)"

    text = re.sub(
        r"urllib\.request\.urlopen\(((?:[^()]|\([^()]*\))*?)\)(?!\s*\))",
        add_urlopen_headers, text)

    # 3. 插入 _UA 定义（放在 import 之后）
    lines = text.splitlines(keepends=True)
    insert_at = 0
    for index, line in enumerate(lines):
        if line.startswith(("import ", "from ")):
            insert_at = index + 1
    lines.insert(insert_at, UA_BLOCK)
    text = "".join(lines)

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

    import ast
    try:
        ast.parse(text)
    except SyntaxError as exc:
        print(f"  {path.name}: 语法错误 {exc}")
        return -1
    print(f"  {path.name}: 改了 {changes} 处")
    return changes


def main() -> int:
    total = 0
    ok = True
    for name in FILES:
        path = TESTS / name
        if not path.is_file():
            print(f"  [跳过] {name} 不存在")
            continue
        result = patch_file(path)
        if result < 0:
            ok = False
        else:
            total += result
    print(f"  合计 {total} 处")
    return 0 if ok else 1


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