#!/usr/bin/env python3
"""
预算文件关键字分析 - 批量解析脚本
解析 /data/www/files/2026-yjs-lzy/ 目录下所有预算文件
匹配17个关键字，提取预算金额，按市州→区县→单位类型分层
输出: /data/www/files/2026-yjs-lzy/parsed_data.json
"""

import os
import re
import json
import subprocess
import tempfile
import shutil
from pathlib import Path

# === 配置 ===
BASE_DIR = "/data/www/files/2026-yjs-lzy"
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):
    """判断是组织部还是党校"""
    name = filename
    if "党校" in name or "_dx_" in name or "-党校" in name or "党校" in name:
        return "党校"
    if "组织部" in name or "_org_" in name or "-组织部" in name or "组织部门" in name or "委员会组织部" in name:
        return "组织部"
    if "经济科技" in name:
        return "其他"
    return "未知"

# 市州判断
def detect_city(filepath, filename):
    """根据文件路径和文件名判断市州"""
    # 先看路径
    parts = filepath.split("/")
    for part in parts:
        for city in CITY_MAP:
            if city in part or part == city:
                return city
        # 也检查路径中的区县名
        for district, city in DISTRICT_TO_CITY.items():
            if district in part:
                return city
    
    # 看文件名中的区县
    for district, city in DISTRICT_TO_CITY.items():
        if district in filename:
            return city
    
    # 看文件名中的市州
    for city in CITY_MAP:
        if city in filename:
            return city
    
    return "未知"

# 区县判断
def detect_district(filepath, filename):
    """根据文件路径和文件名判断区县"""
    full_path = filepath
    
    for district in DISTRICT_TO_CITY:
        if district in full_path:
            return district
    
    # 特殊处理
    if "市中" in full_path and "乐山" in full_path:
        return "市中区"
    
    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:
            return district
    
    # 再匹配文件名
    for district in sorted(DISTRICT_TO_CITY.keys(), key=len, reverse=True):
        if district in filename:
            return district
    
    # 特殊情况
    if "市中" in filename:
        return "市中区"
    
    return "市级"

# === Excel解析 ===
def parse_excel(filepath):
    """解析Excel文件，返回所有文本行和金额信息"""
    import openpyxl
    results = []
    
    try:
        wb = openpyxl.load_workbook(filepath, data_only=True, read_only=True)
    except:
        # 尝试用xlrd处理.xls
        try:
            import xlrd
            wb = xlrd.open_workbook(filepath)
            for sheet_idx in range(wb.nsheets):
                sheet = wb.sheet_by_index(sheet_idx)
                sheet_name = sheet.name
                for row_idx in range(sheet.nrows):
                    row_data = [str(sheet.cell_value(row_idx, col_idx)) for col_idx in range(sheet.ncols)]
                    row_text = " ".join(row_data)
                    if row_text.strip():
                        results.append({
                            "sheet": sheet_name,
                            "row": row_idx + 1,
                            "text": row_text,
                            "cells": row_data,
                            "row_num": row_idx + 1
                        })
            return results
        except Exception as e:
            print(f"  [Excel ERROR] {filepath}: {e}")
            return []
    
    for sheet_name in wb.sheetnames:
        ws = wb[sheet_name]
        for row_idx, row in enumerate(ws.iter_rows(values_only=True)):
            row_data = [str(cell) if cell is not None else "" for cell in row]
            row_text = " ".join(row_data)
            if row_text.strip():
                results.append({
                    "sheet": sheet_name,
                    "row": row_idx + 1,
                    "text": row_text,
                    "cells": row_data,
                    "row_num": row_idx + 1
                })
    
    wb.close()
    return results

def parse_et(filepath):
    """解析.et文件，先转为xlsx再用openpyxl"""
    import openpyxl
    # .et可以直接用openpyxl尝试
    try:
        return parse_excel(filepath)
    except:
        # 转换
        with tempfile.TemporaryDirectory() as tmpdir:
            result = 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):
    """解析.docx文件"""
    from docx import Document
    results = []
    try:
        doc = Document(filepath)
        for para_idx, para in enumerate(doc.paragraphs):
            text = para.text.strip()
            if text:
                results.append({
                    "sheet": "正文",
                    "row": para_idx + 1,
                    "text": text,
                    "cells": [text],
                    "row_num": para_idx + 1
                })
        # 表格
        for table_idx, table in enumerate(doc.tables):
            for row_idx, row in enumerate(table.rows):
                cells = [cell.text.strip() for cell in row.cells]
                row_text = " ".join(cells)
                if row_text.strip():
                    results.append({
                        "sheet": f"表格{table_idx+1}",
                        "row": row_idx + 1,
                        "text": row_text,
                        "cells": cells,
                        "row_num": row_idx + 1
                    })
    except Exception as e:
        print(f"  [DOCX ERROR] {filepath}: {e}")
    return results

def parse_doc(filepath):
    """解析.doc文件，用libreoffice转换"""
    with tempfile.TemporaryDirectory() as tmpdir:
        result = subprocess.run(
            ["libreoffice", "--headless", "--convert-to", "txt", "--outdir", tmpdir, filepath],
            capture_output=True, text=True, timeout=60
        )
        txt_file = os.path.join(tmpdir, Path(filepath).stem + ".txt")
        if os.path.exists(txt_file):
            with open(txt_file, "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],
                        "row_num": idx + 1
                    })
            return results
    return []

# === PDF解析 ===
def parse_pdf(filepath):
    """解析PDF文件，先用pdftotext，失败则OCR"""
    results = []
    
    # 先用pdftotext
    try:
        result = subprocess.run(
            ["pdftotext", "-layout", filepath, "-"],
            capture_output=True, text=True, timeout=120
        )
        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],
                                "row_num": line_in_page
                            })
                else:
                    line = line.strip()
                    if line:
                        line_in_page += 1
                        results.append({
                            "sheet": f"第{page}页",
                            "row": line_in_page,
                            "text": line,
                            "cells": [line],
                            "row_num": line_in_page
                        })
            
            if results:
                return results
    except Exception as e:
        print(f"  [PDF pdftotext ERROR] {filepath}: {e}")
    
    # OCR fallback
    print(f"  [PDF OCR] 尝试OCR: {filepath}")
    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            # 转图片
            subprocess.run(
                ["pdftoppm", "-r", "200", "-png", filepath, os.path.join(tmpdir, "page")],
                capture_output=True, timeout=300
            )
            pages = sorted([f for f in os.listdir(tmpdir) if f.endswith(".png")])
            for page_idx, page_file in enumerate(pages):
                img_path = os.path.join(tmpdir, page_file)
                result = subprocess.run(
                    ["tesseract", img_path, "-", "-l", "chi_sim", "--psm", "6"],
                    capture_output=True, text=True, timeout=120
                )
                lines = result.stdout.split("\n")
                line_in_page = 0
                for line in lines:
                    line = line.strip()
                    if line:
                        line_in_page += 1
                        results.append({
                            "sheet": f"第{page_idx+1}页",
                            "row": line_in_page,
                            "text": line,
                            "cells": [line],
                            "row_num": line_in_page
                        })
    except Exception as e:
        print(f"  [PDF OCR ERROR] {filepath}: {e}")
    
    return results

# === 关键字匹配 ===
def match_keywords(rows, keywords):
    """在解析结果中匹配关键字，返回匹配项列表"""
    matches = []
    for row_data in rows:
        text = row_data["text"]
        for kw in keywords:
            if kw in text:
                # 提取金额
                amount = extract_amount(row_data, text)
                # 提取匹配上下文
                context = extract_context(text, kw)
                
                location = f"{row_data['sheet']}:R{row_data['row']}"
                
                matches.append({
                    "keyword": kw,
                    "context": context,
                    "location": location,
                    "amount": amount,
                    "full_text": text[:200]
                })
    return matches

def extract_amount(row_data, text):
    """从行数据中提取预算金额"""
    # 方法1: 在cells中找金额列
    cells = row_data.get("cells", [])
    
    # 找表头中包含"预算"、"金额"、"总计"、"合计"的列
    # 先尝试从同行cells中找数字
    for cell in cells:
        cell = str(cell).strip()
        # 匹配 XXXX.XX 万元 或 XXXX.XX万
        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("，", ""))
            # 如果大于1000，可能是元，转换为万元
            if val > 1000:
                return round(val / 10000, 2)
            elif val > 0:
                return round(val, 2)
        except:
            pass
    
    # 方法2: 从text中正则提取
    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)
    context = text[start:end].strip()
    if start > 0:
        context = "..." + context
    if end < len(text):
        context = context + "..."
    return context

# === 主流程 ===
def process_all_files():
    all_files = []
    for root, dirs, files in os.walk(BASE_DIR):
        for f in files:
            if f.startswith(".") or f == "Thumbs.db":
                continue
            filepath = os.path.join(root, f)
            rel_path = os.path.relpath(filepath, BASE_DIR)
            ext = os.path.splitext(f)[1].lower()
            
            if ext in ['.xlsx', '.xls', '.et', '.docx', '.doc', '.pdf']:
                all_files.append({
                    "filepath": filepath,
                    "rel_path": rel_path,
                    "filename": f,
                    "ext": ext
                })
    
    print(f"共 {len(all_files)} 个预算文件待解析")
    print("=" * 60)
    
    all_results = []
    
    for i, file_info in enumerate(all_files):
        print(f"\n[{i+1}/{len(all_files)}] 解析: {file_info['rel_path']}")
        
        # 解析文件
        ext = file_info["ext"]
        filepath = file_info["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}")
            rows = []
        
        if not rows:
            print(f"  [无内容]")
            continue
        
        # 匹配关键字
        matches = match_keywords(rows, KEYWORDS)
        
        if not matches:
            print(f"  [无匹配] (共{len(rows)}行)")
            continue
        
        # 判断市州、区县、单位类型
        city = detect_city(file_info["rel_path"], file_info["filename"])
        district = detect_district_name(file_info["rel_path"], file_info["filename"])
        unit_type = detect_unit_type(file_info["filename"])
        
        print(f"  市州={city}, 区县={district}, 单位={unit_type}, 匹配={len(matches)}项")
        
        file_result = {
            "filename": file_info["filename"],
            "rel_path": file_info["rel_path"],
            "city": city,
            "district": district,
            "unit_type": unit_type,
            "ext": ext,
            "matches": matches,
            "total_amount": sum(m["amount"] for m in matches)
        }
        all_results.append(file_result)
    
    print("\n" + "=" * 60)
    print(f"解析完成: {len(all_results)} 个文件有匹配项")
    print(f"总匹配项: {sum(len(r['matches']) for r in all_results)}")
    print(f"总金额: {sum(r['total_amount'] for r in all_results):.2f} 万元")
    
    # 保存JSON
    output_path = "/data/www/files/2026-yjs-lzy/parsed_data.json"
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(all_results, f, ensure_ascii=False, indent=2)
    print(f"\n数据已保存: {output_path}")
    
    return all_results

if __name__ == "__main__":
    process_all_files()