#!/usr/bin/env python3
"""
预算文件关键字分析 - 批量解析脚本 (优化版)
- 跳过OCR，只用pdftotext
- 跳过已解析文件(断点续传)
- 后台运行
"""
import os, re, json, subprocess, tempfile, sys
from pathlib import Path

BASE_DIR = "/data/www/files/2026-yjs-lzy"
OUTPUT_JSON = "/data/www/files/2026-yjs-lzy/parsed_data.json"
PROGRESS_FILE = "/data/www/files/2026-yjs-lzy/parse_progress.json"

KEYWORDS = [
    "网络培训", "培训", "能力提升", "视频拍摄", "党员教育片",
    "课程制作", "直播", "视频", "课件", "系统",
    "信息化", "网络", "数字", "设备", "平台", "维护", "软件"
]

CITY_MAP = {
    "成都": ["双流","大邑","崇州","都江堰","金堂","龙泉驿","青白江","邛崃","郫都","浦江","新都","温江","天府新区"],
    "巴中": ["恩阳","南江","平昌","通江","巴州"],
    "攀枝花": ["东区","仁和","西区","盐边","米易"],
    "绵阳": ["安州","平武","三台","涪城","游仙","梓潼","北川","江油"],
    "乐山": ["井研","峨边","沙湾","沐川","市中","夹江","犍为","金口河","峨眉"],
    "雅安": ["天全","宝兴","汉源","石棉","芦山","荥经","名山","雨城"],
}

DISTRICT_TO_CITY = {}
for city, districts in CITY_MAP.items():
    for d in districts:
        DISTRICT_TO_CITY[d] = city

def detect_unit_type(filename):
    if "党校" in filename or "_dx_" in filename or "-党校" in filename:
        return "党校"
    if "组织部" in filename or "_org_" in filename or "-组织部" in filename or "组织部门" in filename or "委员会组织部" in filename:
        return "组织部"
    if "经济科技" in filename:
        return "其他"
    return "未知"

def detect_city(filepath, filename):
    full = filepath + "/" + filename
    for city in CITY_MAP:
        if city in full:
            return city
    for district, city in DISTRICT_TO_CITY.items():
        if district in full:
            return city
    return "未知"

def detect_district_name(filepath, filename):
    full = filepath + "/" + filename
    for district in sorted(DISTRICT_TO_CITY.keys(), key=len, reverse=True):
        if district in filepath or district in filename:
            return district
    if "市中" in filename:
        return "市中区"
    return "市级"

# === Excel ===
def parse_excel(filepath):
    import openpyxl
    results = []
    try:
        wb = openpyxl.load_workbook(filepath, data_only=True, read_only=True)
    except:
        try:
            import xlrd
            wb = xlrd.open_workbook(filepath)
            for si in range(wb.nsheets):
                sheet = wb.sheet_by_index(si)
                for ri in range(sheet.nrows):
                    row_data = [str(sheet.cell_value(ri, ci)) for ci in range(sheet.ncols)]
                    row_text = " ".join(row_data)
                    if row_text.strip():
                        results.append({"sheet": sheet.name, "row": ri+1, "text": row_text, "cells": row_data})
            return results
        except Exception as e:
            print(f"  [Excel ERROR] {e}", flush=True)
            return []
    for sn in wb.sheetnames:
        ws = wb[sn]
        for ri, row in enumerate(ws.iter_rows(values_only=True)):
            row_data = [str(c) if c is not None else "" for c in row]
            row_text = " ".join(row_data)
            if row_text.strip():
                results.append({"sheet": sn, "row": ri+1, "text": row_text, "cells": row_data})
    wb.close()
    return results

def parse_et(filepath):
    try:
        return parse_excel(filepath)
    except:
        with tempfile.TemporaryDirectory() as tmpdir:
            subprocess.run(["libreoffice","--headless","--convert-to","xlsx","--outdir",tmpdir,filepath],
                          capture_output=True, text=True, timeout=60)
            converted = os.path.join(tmpdir, Path(filepath).stem + ".xlsx")
            if os.path.exists(converted):
                return parse_excel(converted)
        return []

# === Word ===
def parse_docx(filepath):
    from docx import Document
    results = []
    try:
        doc = Document(filepath)
        for pi, para in enumerate(doc.paragraphs):
            t = para.text.strip()
            if t:
                results.append({"sheet":"正文","row":pi+1,"text":t,"cells":[t]})
        for ti, table in enumerate(doc.tables):
            for ri, row in enumerate(table.rows):
                cells = [c.text.strip() for c in row.cells]
                rt = " ".join(cells)
                if rt.strip():
                    results.append({"sheet":f"表格{ti+1}","row":ri+1,"text":rt,"cells":cells})
    except Exception as e:
        print(f"  [DOCX ERROR] {e}", flush=True)
    return results

def parse_doc(filepath):
    with tempfile.TemporaryDirectory() as tmpdir:
        subprocess.run(["libreoffice","--headless","--convert-to","txt","--outdir",tmpdir,filepath],
                      capture_output=True, text=True, timeout=60)
        tf = os.path.join(tmpdir, Path(filepath).stem + ".txt")
        if os.path.exists(tf):
            with open(tf,"r",encoding="utf-8",errors="ignore") as f:
                lines = f.readlines()
            results = []
            for idx, line in enumerate(lines):
                line = line.strip()
                if line:
                    results.append({"sheet":"正文","row":idx+1,"text":line,"cells":[line]})
            return results
    return []

# === PDF (only pdftotext, skip OCR) ===
def parse_pdf(filepath):
    results = []
    try:
        result = subprocess.run(["pdftotext","-layout",filepath,"-"],
                              capture_output=True, text=True, timeout=60)
        if result.returncode == 0 and result.stdout.strip():
            lines = result.stdout.split("\n")
            page = 1
            line_in_page = 0
            for line in lines:
                if "\f" in line:
                    parts = line.split("\f")
                    for i, part in enumerate(parts):
                        if i > 0:
                            page += 1
                            line_in_page = 0
                        part = part.strip()
                        if part:
                            line_in_page += 1
                            results.append({"sheet":f"第{page}页","row":line_in_page,"text":part,"cells":[part]})
                else:
                    line = line.strip()
                    if line:
                        line_in_page += 1
                        results.append({"sheet":f"第{page}页","row":line_in_page,"text":line,"cells":[line]})
        if results:
            return results
    except Exception as e:
        print(f"  [PDF ERROR] {e}", flush=True)
    # 跳过OCR，直接返回空
    print(f"  [PDF无文本层，跳过]", flush=True)
    return []

# === 关键字匹配 ===
def match_keywords(rows, keywords):
    matches = []
    for rd in rows:
        text = rd["text"]
        for kw in keywords:
            if kw in text:
                amount = extract_amount(rd, text)
                context = extract_context(text, kw)
                location = f"{rd['sheet']}:R{rd['row']}"
                matches.append({
                    "keyword": kw,
                    "context": context,
                    "location": location,
                    "amount": amount,
                    "full_text": text[:200]
                })
    return matches

def extract_amount(row_data, text):
    cells = row_data.get("cells", [])
    for cell in cells:
        cell = str(cell).strip()
        m = re.search(r'(\d+\.?\d*)\s*万元?', cell)
        if m:
            val = float(m.group(1))
            if 0 < val < 100000:
                return round(val, 2)
        try:
            val = float(cell.replace(",","").replace("，",""))
            if val > 1000:
                return round(val / 10000, 2)
            elif val > 0:
                return round(val, 2)
        except:
            pass
    m = re.search(r'(\d+\.?\d*)\s*万元?', text)
    if m:
        val = float(m.group(1))
        if 0 < val < 100000:
            return round(val, 2)
    return 0.0

def extract_context(text, keyword):
    idx = text.find(keyword)
    if idx == -1:
        return text[:80]
    start = max(0, idx - 20)
    end = min(len(text), idx + len(keyword) + 30)
    ctx = text[start:end].strip()
    if start > 0: ctx = "..." + ctx
    if end < len(text): ctx = ctx + "..."
    return ctx

# === 主流程 ===
def main():
    # 收集文件
    all_files = []
    for root, dirs, files in os.walk(BASE_DIR):
        for f in files:
            if f.startswith(".") or f == "Thumbs.db" or f == "parse_budget.py" or f == "parse_budget_fast.py":
                continue
            ext = os.path.splitext(f)[1].lower()
            if ext in ['.xlsx','.xls','.et','.docx','.doc','.pdf']:
                all_files.append({
                    "filepath": os.path.join(root, f),
                    "rel_path": os.path.relpath(os.path.join(root, f), BASE_DIR),
                    "filename": f,
                    "ext": ext
                })
    
    # 断点续传：加载已解析的
    done_files = set()
    all_results = []
    if os.path.exists(PROGRESS_FILE):
        with open(PROGRESS_FILE, "r", encoding="utf-8") as f:
            progress = json.load(f)
            done_files = set(progress.get("done_files", []))
            all_results = progress.get("results", [])
            print(f"断点续传: 已完成 {len(done_files)} 个文件", flush=True)
    
    print(f"共 {len(all_files)} 个预算文件，已解析 {len(done_files)}，剩余 {len(all_files) - len(done_files)}", flush=True)
    print("=" * 60, flush=True)
    
    for i, fi in enumerate(all_files):
        if fi["rel_path"] in done_files:
            continue
        
        print(f"\n[{i+1}/{len(all_files)}] {fi['rel_path']}", flush=True)
        
        ext = fi["ext"]
        filepath = fi["filepath"]
        
        try:
            if ext in ['.xlsx','.xls']:
                rows = parse_excel(filepath)
            elif ext == '.et':
                rows = parse_et(filepath)
            elif ext == '.docx':
                rows = parse_docx(filepath)
            elif ext == '.doc':
                rows = parse_doc(filepath)
            elif ext == '.pdf':
                rows = parse_pdf(filepath)
            else:
                continue
        except Exception as e:
            print(f"  [解析失败] {e}", flush=True)
            rows = []
        
        if not rows:
            print(f"  [无内容]", flush=True)
            done_files.add(fi["rel_path"])
            continue
        
        matches = match_keywords(rows, KEYWORDS)
        
        if not matches:
            print(f"  [无匹配] (共{len(rows)}行)", flush=True)
            done_files.add(fi["rel_path"])
            continue
        
        city = detect_city(fi["rel_path"], fi["filename"])
        district = detect_district_name(fi["rel_path"], fi["filename"])
        unit_type = detect_unit_type(fi["filename"])
        
        print(f"  {city}/{district}/{unit_type} 匹配={len(matches)}项 金额={sum(m['amount'] for m in matches):.2f}万", flush=True)
        
        file_result = {
            "filename": fi["filename"],
            "rel_path": fi["rel_path"],
            "city": city,
            "district": district,
            "unit_type": unit_type,
            "ext": ext,
            "matches": matches,
            "total_amount": round(sum(m["amount"] for m in matches), 2)
        }
        all_results.append(file_result)
        done_files.add(fi["rel_path"])
        
        # 每5个文件保存一次进度
        if (i + 1) % 5 == 0:
            with open(PROGRESS_FILE, "w", encoding="utf-8") as f:
                json.dump({"done_files": list(done_files), "results": all_results}, f, ensure_ascii=False)
    
    # 最终保存
    print("\n" + "=" * 60, flush=True)
    print(f"解析完成: {len(all_results)} 个文件有匹配项", flush=True)
    print(f"总匹配项: {sum(len(r['matches']) for r in all_results)}", flush=True)
    print(f"总金额: {sum(r['total_amount'] for r in all_results):.2f} 万元", flush=True)
    
    with open(OUTPUT_JSON, "w", encoding="utf-8") as f:
        json.dump(all_results, f, ensure_ascii=False, indent=2)
    print(f"\n数据已保存: {OUTPUT_JSON}", flush=True)
    
    # 清理进度文件
    if os.path.exists(PROGRESS_FILE):
        os.remove(PROGRESS_FILE)

if __name__ == "__main__":
    main()