#!/usr/bin/env python3
"""补提取已有记录的附件正文"""
import json, re, sys, time, os, sqlite3, subprocess, tempfile
from urllib.parse import urljoin, unquote
import requests
from bs4 import BeautifulSoup

BASE_URL = "https://www.zjcx.gov.cn"
SITE_NAME = "长兴县人民政府-乡镇部门非规范性文件"
DB_PATH = "/root/search.db"
ATTACH_DIR = "/tmp/zjcx_attachments"

HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
    "Accept-Language": "zh-CN,zh;q=0.9",
}
session = requests.Session()
session.headers.update(HEADERS)
os.makedirs(ATTACH_DIR, exist_ok=True)


def extract_text_from_file(path):
    ext = os.path.splitext(path)[1].lower()
    text = ""
    if ext == '.docx':
        try:
            import docx
            doc = docx.Document(path)
            text = '\n'.join(p.text for p in doc.paragraphs if p.text.strip())
        except:
            pass
    elif ext == '.doc':
        try:
            result = subprocess.run(['catdoc', path], capture_output=True, text=True, timeout=30)
            text = result.stdout.strip()
        except:
            try:
                result = subprocess.run(['antiword', path], capture_output=True, text=True, timeout=30)
                text = result.stdout.strip()
            except:
                pass
    elif ext in ('.xlsx', '.xls'):
        try:
            import openpyxl
            wb = openpyxl.load_workbook(path, read_only=True, data_only=True)
            texts = []
            for ws in wb.worksheets:
                for row in ws.iter_rows():
                    for cell in row:
                        if cell.value is not None:
                            texts.append(str(cell.value))
            text = '\n'.join(texts)
        except:
            pass
    elif ext == '.pdf':
        try:
            result = subprocess.run(['pdftotext', path, '-'], capture_output=True, text=True, timeout=30)
            text = result.stdout.strip()
        except:
            pass
    return text


def download_attachment(url):
    try:
        r = session.get(url, timeout=60, verify=False, stream=True)
        if r.status_code != 200:
            return None
        url_path = unquote(url.split('?')[0].split('/')[-1])
        if not url_path:
            url_path = f"attach_{int(time.time())}"
        local_path = os.path.join(ATTACH_DIR, url_path)
        with open(local_path, 'wb') as f:
            for chunk in r.iter_content(8192):
                f.write(chunk)
        return local_path
    except:
        return None


def get_attachments_from_detail(detail_url):
    """重新抓取详情页，提取所有附件URL"""
    attachments = []
    try:
        resp = session.get(detail_url, timeout=60, verify=False)
        resp.encoding = "utf-8"
        soup = BeautifulSoup(resp.text, "html.parser")
        zw = soup.find("div", class_="zhengw")
        if zw:
            # script-based
            for script in zw.find_all("script"):
                for m in re.finditer(r'pdfurl\s*=\s*"([^"]+)"', script.string or ""):
                    full_attach = urljoin(BASE_URL, m.group(1))
                    attachments.append(full_attach)
            # a-tag based
            for a_tag in zw.find_all("a", href=True):
                href = a_tag["href"]
                if "document/download" in href or re.search(r'\.(docx?|xlsx?|pdf)$', href, re.I):
                    full_url = urljoin(BASE_URL, href)
                    if full_url not in attachments:
                        attachments.append(full_url)
    except:
        pass
    return attachments


def main():
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()

    # 找到需要补提取的记录
    cursor.execute("""
        SELECT id, page_url, summary FROM gov_raw 
        WHERE site_name = ? AND (content IS NULL OR content = '' OR length(content) < 100) 
        AND summary != ''
    """, (SITE_NAME,))
    rows = cursor.fetchall()
    print(f"需要补提取: {len(rows)} 条")

    fixed = 0
    failed = 0

    for row_id, page_url, summary in rows:
        print(f"\n[{fixed+1}/{len(rows)}] id={row_id}")
        attachments = get_attachments_from_detail(page_url)
        if not attachments:
            print(f"  未找到附件链接，跳过")
            failed += 1
            continue

        print(f"  找到 {len(attachments)} 个附件")
        attach_text = ""
        for au in attachments:
            print(f"    下载: {au[:80]}...")
            local_path = download_attachment(au)
            if local_path:
                text = extract_text_from_file(local_path)
                if text:
                    attach_text += text + "\n\n"
                    print(f"    提取: {len(text)} 字符")
                try:
                    os.remove(local_path)
                except:
                    pass

        if attach_text:
            content_html = "<div class=\"zhengw\"><p>" + attach_text.replace("\n", "</p><p>") + "</p></div>"
            cursor.execute("UPDATE gov_raw SET content = ? WHERE id = ?", (content_html, row_id))
            fixed += 1
            print(f"  ✓ 已更新 {len(attach_text)} 字")
        else:
            print(f"  ⚠ 未能提取任何文本")
            failed += 1

        if fixed % 20 == 0:
            conn.commit()
            print(f"  --- 已提交 {fixed} 条 ---")

        time.sleep(0.5)

    conn.commit()
    conn.close()

    print(f"\n=== 完成 ===")
    print(f"修复: {fixed} 条")
    print(f"失败: {failed} 条")


if __name__ == "__main__":
    main()