#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""批量修「手动写 gov_search → IntegrityError → 连带回滚 gov_raw」的脚本（增强版）。

改进点：
  · 不依赖固定变量名/固定格式 —— 扫到含 `INTO gov_search` 的语句，向上找它所属的
    execute 调用行（可能是多行），在其前面按同缩进插入 `<连接变量>.commit()`。
  · 连接变量从脚本里 `xxx = sqlite3.connect(` 自动探测，探测不到则跳过（交人工）。
  · 已含修复标记的跳过（幂等）。
"""
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
MARK = "先提交 gov_raw"
ONLY = None
for a in sys.argv[1:]:
    if a.startswith("--list="):
        ONLY = set(l.strip() for l in io.open(a.split("=", 1)[1], encoding="utf-8") if l.strip())

CONN = re.compile(r"^\s*([A-Za-z_]\w*)\s*=\s*sqlite3\.connect\(", re.M)
EXEC_START = re.compile(r"^\s*([A-Za-z_][\w\.]*)\.(execute|executemany)\(", re.M)

fixed, skipped, no_conn = [], [], []
for fn in sorted(os.listdir(D)):
    if not fn.startswith("crawl_") or not fn.endswith(".py"):
        continue
    if ONLY is not None and fn not in ONLY:
        continue
    p = os.path.join(D, fn)
    src = io.open(p, encoding="utf-8", errors="ignore").read()
    if "INTO gov_search" not in src or MARK in src:
        continue
    m = CONN.search(src)
    if not m:
        no_conn.append(fn)
        continue
    connvar = m.group(1)

    lines = src.split("\n")
    # 找出所有含 INTO gov_search 的行号
    hit_lines = [i for i, l in enumerate(lines) if "INTO gov_search" in l]
    if not hit_lines:
        continue
    inserted = 0
    out = list(lines)
    # 从后往前插，避免行号漂移
    done_rows = set()
    for hl in reversed(hit_lines):
        # 向上找最近的 execute 起始行
        j = hl
        while j >= 0 and not EXEC_START.match(out[j]):
            j -= 1
        if j < 0 or j in done_rows:
            continue
        done_rows.add(j)
        ind = re.match(r"^(\s*)", out[j]).group(1)
        if MARK in "\n".join(out[max(0, j - 3):j]):
            continue
        out[j:j] = [
            ind + "# 2026-09-22: 先提交 gov_raw —— 库上触发器已维护 FTS，下面这条手动写入会",
            ind + "#   因 rowid 重复而 IntegrityError；不先 commit 会把 gov_raw 那条一并回滚（静默丢数据）",
            ind + "%s.commit()" % connvar,
        ]
        inserted += 1
    if not inserted:
        skipped.append(fn)
        continue
    new = "\n".join(out)
    try:
        ast.parse(new)
    except SyntaxError as e:
        skipped.append("%s(%s)" % (fn, str(e)[:40]))
        continue
    if not DRY:
        os.makedirs(ARCH, exist_ok=True)
        bak = os.path.join(ARCH, fn + ".bak_20260922_precommit")
        if not os.path.exists(bak):
            shutil.copy2(p, bak)
        io.open(p, "w", encoding="utf-8").write(new)
    fixed.append(fn)
    print("  ✅ %-44s [%s] 插 %d 处" % (fn, connvar, inserted))

print()
print("已修 %d 个%s" % (len(fixed), "  [DRY-RUN]" if DRY else ""))
if skipped:
    print("跳过（语法失败/无插入点）%d 个: %s" % (len(skipped), ", ".join(skipped[:6])))
if no_conn:
    print("⚠️ 探不到 sqlite3.connect 变量名 %d 个: %s" % (len(no_conn), ", ".join(no_conn[:8])))