#!/usr/bin/env python3
"""OCR处理11个扫描版PDF"""
import os, sys, json, re, subprocess
from pathlib import Path

BASE = Path("/data/www/files/2026-yjs-lzy")

KEYWORDS = [
    "党员教育", "党员培训", "干部教育", "干部培训",
    "党员管理", "党员干部", "党建", "基层党组织",
    "党校", "远程教育", "网络教育", "信息化",
    "人才", "驻村", "乡村振兴", "主题教育",
    "党代表"
]

def find_keywords_in_text(text, source_info):
    results = []
    for kw in KEYWORDS:
        if kw in text:
            idx = 0
            while True:
                pos = text.find(kw, idx)
                if pos == -1:
                    break
                start = max(0, pos - 100)
                end = min(len(text), pos + len(kw) + 200)
                context = text[start:end].replace('\n', ' ').strip()
                # OCR文本中提取金额
                amounts = []
                for pat in [r'(\d+(?:\.\d+)?)\s*万元', r'(\d+(?:\.\d+)?)\s*万\s*元']:
                    for m in re.finditer(pat, context, re.IGNORECASE):
                        try:
                            val = float(m.group(1))
                            if 0.01 <= val <= 50000:
                                amounts.append(val)
                        except:
                            pass
                amount = max(amounts) if amounts else 0
                results.append({
                    "keyword": kw, "amount": amount,
                    "source": source_info, "context": context[:300]
                })
                idx = pos + len(kw)
    return results

with open(BASE / "ocr_pending.json", "r", encoding="utf-8") as f:
    ocr_files = json.load(f)

print(f"共{len(ocr_files)}个文件需OCR")

ocr_results = {}
for i, item in enumerate(ocr_files):
    rel_path = item["path"]
    filepath = str(BASE / rel_path)
    print(f"[{i+1}/{len(ocr_files)}] {rel_path}", end="", flush=True)
    
    try:
        # 先用pdftoppm转图片，再用tesseract OCR
        # 降低DPI到150加快速度
        import tempfile
        with tempfile.TemporaryDirectory() as tmpdir:
            # 转第一页为图片（预算文件通常关键信息在前几页）
            out = subprocess.run(
                ["pdftoppm", "-r", "150", "-f", "1", "-l", "10", filepath, 
                 os.path.join(tmpdir, "page")],
                capture_output=True, text=True, timeout=60
            )
            if out.returncode != 0:
                print(f" -> pdftoppm失败")
                ocr_results[rel_path] = {
                    "city": item["city"], "unit_type": item["unit_type"],
                    "file_type": "PDF-OCR失败", "matches": []
                }
                continue
            
            # 找所有生成的图片
            images = sorted([f for f in os.listdir(tmpdir) if f.endswith('.ppm') or f.endswith('.png')])
            full_text = ""
            for img in images:
                img_path = os.path.join(tmpdir, img)
                ocr_out = subprocess.run(
                    ["tesseract", img_path, "-", "-l", "chi_sim", "--psm", "6"],
                    capture_output=True, text=True, timeout=60
                )
                if ocr_out.returncode == 0:
                    full_text += ocr_out.stdout + "\n"
            
            matches = find_keywords_in_text(full_text, "PDF-OCR")
            ocr_results[rel_path] = {
                "city": item["city"], "unit_type": item["unit_type"],
                "file_type": "PDF-OCR", "matches": matches
            }
            print(f" -> {len(matches)}匹配")
    except Exception as e:
        print(f" -> 错误: {str(e)[:80]}")
        ocr_results[rel_path] = {
            "city": item["city"], "unit_type": item["unit_type"],
            "file_type": "PDF-OCR错误", "matches": [], "error": str(e)[:200]
        }

# 保存OCR结果
with open(BASE / "parsed_ocr.json", "w", encoding="utf-8") as f:
    json.dump(ocr_results, f, ensure_ascii=False, indent=2)

# 合并到fast结果
with open(BASE / "parsed_fast.json", "r", encoding="utf-8") as f:
    fast = json.load(f)

for k, v in ocr_results.items():
    fast[k] = v

with open(BASE / "parsed_final.json", "w", encoding="utf-8") as f:
    json.dump(fast, f, ensure_ascii=False, indent=2)

total_matches = sum(len(d["matches"]) for d in fast.values())
total_amount = sum(m["amount"] for d in fast.values() for m in d["matches"] if m.get("amount", 0) > 0)
print(f"\n最终合并: {len(fast)}文件, {total_matches}匹配, {total_amount:.2f}万元")
print("已保存: parsed_final.json")