#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""给缺 timeout= 的 requests 调用补上超时（防连接停顿 → 永久挂住 → 600s 被杀）。"""
import ast
import io
import os
import re
import shutil
import sys

D = "/root/gov_crawler"
ARCH = os.path.join(D, "Archive")
DRY = "--dry-run" in sys.argv
TO = 30

TARGETS = ["crawl_lankao_tzgg.py", "crawl_yc_hb.py", "list_no_data.py",
           "crawl_fx.py", "spider.py", "spider_unified.py",
           "crawl_jxyy_tzgg.py", "crawl_liyang_tzgg.py",
           "crawl_liyang_xxgk.py", "crawl_zdhjkj_notice.py"]

pat = re.compile(r"(requests\.(?:get|post|head|put)|\b(?:session|http|s|sess|req)\.(?:get|post|head|put))\s*\(", re.I)

for fn in TARGETS:
    p = os.path.join(D, fn)
    if not os.path.exists(p):
        print("   ⚠️  %s 不存在" % fn)
        continue
    src = io.open(p, encoding="utf-8", errors="ignore").read()
    edits = []
    for m in pat.finditer(src):
        i = m.end() - 1
        depth = 0
        j = i
        while j < len(src):
            if src[j] == "(":
                depth += 1
            elif src[j] == ")":
                depth -= 1
                if depth == 0:
                    break
            j += 1
        seg = src[i:j]
        if "timeout" in seg:
            continue
        edits.append(j)   # 在闭括号前插入

    if not edits:
        print("   ⏭  %s 无需修改" % fn)
        continue

    new = src
    for pos in sorted(edits, reverse=True):
        head = new[:pos].rstrip()
        if head.endswith(","):
            new = new[:pos].rstrip() + " timeout=%d," % TO + new[pos:]
        else:
            new = new[:pos].rstrip() + ", timeout=%d" % TO + new[pos:]

    try:
        ast.parse(new)
    except SyntaxError as e:
        print("   ❌ %s 语法失败: %s" % (fn, str(e)[:70]))
        continue

    if not DRY:
        os.makedirs(ARCH, exist_ok=True)
        bak = os.path.join(ARCH, fn + ".bak_20260921_reqtimeout")
        if not os.path.exists(bak):
            shutil.copy2(p, bak)
        io.open(p, "w", encoding="utf-8").write(new)
    print("   ✅ %-34s 补了 %d 处 timeout=%ds" % (fn, len(edits), TO))
print()
print("完成%s" % ("  [DRY-RUN]" if DRY else ""))