#!/usr/bin/env python3
"""预算文件分级解析脚本 v2
- Excel/Word: 秒级解析
- PDF有文本层: pdftotext秒级提取
- PDF无文本层: 标记为needs_ocr，后面单独处理
"""
import os, sys, json, re, subprocess, traceback
from pathlib import Path

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

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

def classify_file(filepath):
    """判断单位类型：组织部 or 党校"""
    name = os.path.basename(filepath)
    parent = os.path.basename(os.path.dirname(filepath))
    combined = name + " " + parent
    if "组织部" in combined or "组织部门" in combined or "组织委员会" in combined:
        return "组织部"
    elif "党校" in combined or "行政学校" in combined:
        return "党校"
    elif "经济科技" in combined:
        return "其他"
    else:
        return "其他"

def classify_city(filepath):
    """判断市州"""
    name = os.path.basename(filepath)
    parent = os.path.basename(os.path.dirname(filepath))
    combined = name + " " + parent
    
    city_map = {
        "成都": ["双流", "大邑", "崇州", "都江堰", "金堂", "龙泉驿", "青白江", "邛崃", "郫都", "浦江", "新都"],
        "巴中": ["巴中", "恩阳", "南江", "平昌", "通江"],
        "攀枝花": ["攀枝花", "东区", "仁和", "西区", "盐边", "米易"],
        "绵阳": ["绵阳", "安州", "平武", "三台", "涪城", "游仙", "梓潼", "北川", "江油"],
        "乐山": ["乐山", "井研", "峨边", "沙湾", "沐川", "市中区", "夹江", "犍为", "金口河", "峨眉山"],
        "雅安": ["雅安", "天全", "宝兴", "汉源", "石棉", "芦山", "荥经", "名山", "雨城"],
    }
    for city, districts in city_map.items():
        for d in districts:
            if d in combined:
                return city
    return "其他"

def extract_amount(text):
    """从文本中提取金额（万元），限制合理范围"""
    amounts = []
    # 匹配 "XX万元" "XX.XX万元" "合计XX万元" 等
    patterns = [
        r'(\d+(?:\.\d+)?)\s*万元',
        r'(\d+(?:\.\d+)?)\s*万\s*元',
        r'金额[：:]\s*(\d+(?:\.\d+)?)',
        r'合计[：:]\s*(\d+(?:\.\d+)?)',
        r'预算[：:]\s*(\d+(?:\.\d+)?)',
    ]
    for pat in patterns:
        for m in re.finditer(pat, text, re.IGNORECASE):
            try:
                val = float(m.group(1))
                # 合理范围：0.01万元 ~ 5000万元
                if 0.01 <= val <= 50000:
                    amounts.append(val)
            except:
                pass
    return amounts

def extract_amount_from_number(val_str):
    """从纯数字字符串提取金额，排除编码/编号"""
    if not val_str:
        return None
    val_str = str(val_str).strip()
    # 排除：纯编码（含字母）、超长数字串（科目编码如51111122T000004930678）
    if re.search(r'[a-zA-Z]', val_str):
        return None
    if len(val_str) > 12:
        return None
    # 排除年份
    if val_str in ['2026', '2025', '2024', '2023', '2022', '2021', '2020']:
        return None
    try:
        val = float(val_str)
        # 合理范围：0.01 ~ 50000万元
        if 0.01 <= val <= 50000:
            return val
    except:
        pass
    return None

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
                # 取上下文（前后100字符）
                start = max(0, pos - 100)
                end = min(len(text), pos + len(kw) + 200)
                context = text[start:end].replace('\n', ' ').strip()
                # 尝试从上下文提取金额
                amounts = extract_amount(context)
                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

def parse_excel(filepath):
    """解析Excel文件"""
    results = []
    try:
        from openpyxl import load_workbook
        wb = load_workbook(filepath, data_only=True, read_only=False)
        for ws in wb.worksheets:
            for row_idx, row in enumerate(ws.iter_rows(values_only=False), 1):
                for cell in row:
                    if cell.value is None:
                        continue
                    cell_str = str(cell.value)
                    col_letter = cell.column_letter if hasattr(cell, 'column_letter') else ''
                    source = f"Sheet:{ws.title} 行{row_idx} 列{col_letter}"
                    for kw in KEYWORDS:
                        if kw in cell_str:
                            # 取同行其他单元格的数字作为可能的金额
                            row_amounts = []
                            for c2 in row:
                                if c2.value is not None and c2 != cell:
                                    amt = extract_amount_from_number(str(c2.value))
                                    if amt:
                                        row_amounts.append(amt)
                            amount = max(row_amounts) if row_amounts else 0
                            results.append({
                                "keyword": kw,
                                "amount": amount,
                                "source": source,
                                "context": cell_str[:200]
                            })
        try:
            wb.close()
        except:
            pass
    except Exception as e:
        # 尝试xlrd for .xls
        try:
            import xlrd
            wb = xlrd.open_workbook(filepath)
            for ws in wb.sheets():
                for row_idx in range(ws.nrows):
                    row_vals = [ws.cell_value(row_idx, c) for c in range(ws.ncols)]
                    for col_idx, val in enumerate(row_vals):
                        if val is None:
                            continue
                        cell_str = str(val)
                        source = f"Sheet:{ws.name} 行{row_idx+1} 列{col_idx+1}"
                        for kw in KEYWORDS:
                            if kw in cell_str:
                                row_amounts = []
                                for v2 in row_vals:
                                    if v2 is not None and str(v2) != cell_str:
                                        amt = extract_amount_from_number(str(v2))
                                        if amt:
                                            row_amounts.append(amt)
                                amount = max(row_amounts) if row_amounts else 0
                                results.append({
                                    "keyword": kw,
                                    "amount": amount,
                                    "source": source,
                                    "context": cell_str[:200]
                                })
        except:
            results.append({"keyword": "解析失败", "amount": 0, "source": f"错误: {str(e)[:100]}", "context": ""})
    return results

def parse_word(filepath):
    """解析Word文件"""
    results = []
    text = ""
    try:
        if filepath.endswith('.docx'):
            from docx import Document
            doc = Document(filepath)
            for para in doc.paragraphs:
                text += para.text + "\n"
            for table in doc.tables:
                for row in table.rows:
                    row_text = []
                    for cell in row.cells:
                        row_text.append(cell.text)
                    text += " | ".join(row_text) + "\n"
        elif filepath.endswith('.doc'):
            # 用antiword或catdoc
            try:
                out = subprocess.run(["antiword", filepath], capture_output=True, text=True, timeout=30)
                if out.returncode == 0:
                    text = out.stdout
                else:
                    out2 = subprocess.run(["catdoc", filepath], capture_output=True, text=True, timeout=30)
                    text = out2.stdout if out2.returncode == 0 else ""
            except:
                text = ""
        results = find_keywords_in_text(text, f"Word文档")
    except Exception as e:
        results.append({"keyword": "解析失败", "amount": 0, "source": f"错误: {str(e)[:100]}", "context": ""})
    return results

def parse_pdf_text(filepath):
    """用pdftotext提取PDF文本层"""
    try:
        out = subprocess.run(
            ["pdftotext", "-layout", filepath, "-"],
            capture_output=True, text=True, timeout=30
        )
        if out.returncode == 0 and len(out.stdout.strip()) > 50:
            return out.stdout, False  # False = 不需要OCR
        return None, True  # 需要OCR
    except subprocess.TimeoutExpired:
        return None, True
    except:
        return None, True

def parse_pdf(filepath):
    """解析PDF：先pdftotext，需要OCR则标记"""
    text, needs_ocr = parse_pdf_text(filepath)
    if needs_ocr or text is None:
        return [], True  # 返回空结果，标记需要OCR
    results = find_keywords_in_text(text, "PDF文本")
    return results, False

def main():
    # 收集所有预算文件（排除.py/.json/.html/.js/.txt等）
    valid_exts = {'.xlsx', '.xls', '.et', '.pdf', '.doc', '.docx'}
    all_files = []
    for root, dirs, files in os.walk(BASE):
        for f in files:
            fp = os.path.join(root, f)
            ext = os.path.splitext(f)[1].lower()
            if ext in valid_exts:
                all_files.append(fp)
    all_files.sort()
    
    print(f"共{len(all_files)}个预算文件")
    
    fast_results = {}  # 非OCR结果
    ocr_files = []     # 需要OCR的文件
    
    for i, filepath in enumerate(all_files):
        rel_path = os.path.relpath(filepath, BASE)
        ext = os.path.splitext(filepath)[1].lower()
        city = classify_city(filepath)
        unit_type = classify_file(filepath)
        
        print(f"[{i+1}/{len(all_files)}] {rel_path}", end="", flush=True)
        
        try:
            if ext in ('.xlsx', '.xls', '.et'):
                matches = parse_excel(filepath)
                fast_results[rel_path] = {
                    "city": city, "unit_type": unit_type,
                    "file_type": "Excel", "matches": matches
                }
                print(f" -> {len(matches)}匹配")
            elif ext in ('.doc', '.docx'):
                matches = parse_word(filepath)
                fast_results[rel_path] = {
                    "city": city, "unit_type": unit_type,
                    "file_type": "Word", "matches": matches
                }
                print(f" -> {len(matches)}匹配")
            elif ext == '.pdf':
                matches, needs_ocr = parse_pdf(filepath)
                if needs_ocr:
                    ocr_files.append({"path": rel_path, "city": city, "unit_type": unit_type})
                    print(f" -> 需OCR")
                else:
                    fast_results[rel_path] = {
                        "city": city, "unit_type": unit_type,
                        "file_type": "PDF文本", "matches": matches
                    }
                    print(f" -> {len(matches)}匹配")
        except Exception as e:
            print(f" -> 错误: {str(e)[:80]}")
            fast_results[rel_path] = {
                "city": city, "unit_type": unit_type,
                "file_type": ext, "matches": [],
                "error": str(e)[:200]
            }
    
    # 统计
    total_matches = sum(len(d["matches"]) for d in fast_results.values())
    total_amount = sum(
        m["amount"] for d in fast_results.values() 
        for m in d["matches"] if m.get("amount", 0) > 0
    )
    
    print(f"\n{'='*60}")
    print(f"非OCR解析完成:")
    print(f"  文件数: {len(fast_results)}")
    print(f"  匹配项: {total_matches}")
    print(f"  总金额: {total_amount:.2f}万元")
    print(f"  需OCR文件: {len(ocr_files)}个")
    for f in ocr_files:
        print(f"    - {f['path']} [{f['city']}/{f['unit_type']}]")
    
    # 保存结果
    with open(BASE / "parsed_fast.json", "w", encoding="utf-8") as f:
        json.dump(fast_results, f, ensure_ascii=False, indent=2)
    with open(BASE / "ocr_pending.json", "w", encoding="utf-8") as f:
        json.dump(ocr_files, f, ensure_ascii=False, indent=2)
    
    print(f"\n结果已保存: parsed_fast.json + ocr_pending.json")

if __name__ == "__main__":
    main()