#!/usr/bin/env python3
"""为缺 BeautifulSoup 的脚本补 import + 统一 --pages 参数兼容
1. 检查/补 from bs4 import BeautifulSoup
2. 检查 __main__ 参数解析: 支持 --pages=N 和位置参数
"""
import os, re, glob, ast

DIR = '/root/gov_crawler'

def ensure_bs_import(s):
    if 'from bs4 import BeautifulSoup' in s or 'import BeautifulSoup' in s:
        return s, False
    # 在 import 块后插入
    m = re.search(r'^import .*$', s, re.M)
    if m:
        pos = m.end()
        # 找到 import 块末尾（连续 import/from 行）
        lines = s.split('\n')
        # 简单: 在第一个 import 行前插入
        s = s[:pos] + '\nfrom bs4 import BeautifulSoup' + s[pos:]
        return s, True
    return s, False

def ensure_pages_compat(s):
    """把 'int(sys.argv[1])' 风格改为支持 --pages=N"""
    if '--pages' in s:
        return s, False
    # 常见模式: pages = int(sys.argv[1]) if len(sys.argv) > 1 else N
    m = re.search(r"pages\s*=\s*int\(sys\.argv\[1\]\)\s*if\s*len\(sys\.argv\)\s*>\s*1\s*else\s*(\d+)", s)
    if m:
        default = m.group(1)
        new = f'''def _get_pages():
    import sys as _s
    for i, a in enumerate(_s.argv):
        if a.startswith("--pages="):
            try:
                return int(a.split("=", 1)[1])
            except ValueError:
                pass
        elif a.isdigit():
            return int(a)
    return {default}
pages = _get_pages()'''
        s = s.replace(m.group(0), new)
        return s, True
    return s, False

fixed_imp, fixed_pg = [], []
for p in glob.glob(os.path.join(DIR, 'crawl_*.py')):
    s = open(p).read()
    if 'def html_table_to_html' not in s and 'def html_table_to_markdown' not in s:
        continue
    changed = False
    s, c1 = ensure_bs_import(s)
    if c1: changed = True
    s, c2 = ensure_pages_compat(s)
    if c2: changed = True
    if changed:
        try:
            ast.parse(s)
            open(p, 'w').write(s)
            name = os.path.basename(p)
            if c1: fixed_imp.append(name)
            if c2: fixed_pg.append(name)
        except Exception as e:
            print(f"FAIL {os.path.basename(p)}: {e}")

print(f"补 import ({len(fixed_imp)}):")
for f in fixed_imp: print('  ', f)
print(f"补 pages 兼容 ({len(fixed_pg)}):")
for f in fixed_pg: print('  ', f)