#!/usr/bin/env python3
"""
预算文件关键字分析脚本
解析Excel/Word/PDF预算文件，匹配17个关键字，提取预算金额
按市州→区县→组织部/党校分层整理
"""
import os
import re
import json
import subprocess
import tempfile
from pathlib import Path

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

BASE_DIR = "/data/www/files/2026-yjs-cyl"

def classify_file(filepath):
    """根据文件路径判断市州、区县、单位类型"""
    rel = os.path.relpath(filepath, BASE_DIR)
    parts = rel.split(os.sep)
    
    city = ""
    district = ""
    unit_type = ""  # 组织部 / 党校
    
    fname = os.path.basename(filepath)
    
    # 从文件名或路径判断市州
    if "内江" in fname or "neijiang" in fname.lower() or (len(parts) > 1 and "内江" in parts[0]):
        city = "内江市"
    elif "自贡" in fname or "zigong" in fname.lower() or (len(parts) > 1 and "自贡" in parts[0]):
        city = "自贡市"
    elif "东部新区" in fname:
        city = "成都市"
        district = "东部新区"
    
    # 从文件名判断区县
    district_patterns = [
        ("龙泉驿", "龙泉驿区"),
        ("富顺", "富顺县"),
        ("荣县", "荣县"),
        ("沿滩", "沿滩区"),
        ("大安", "大安区"),
        ("贡井", "贡井区"),
        ("自流井", "自流井区"),
        ("隆昌", "隆昌市"),
        ("东兴", "东兴区"),
        ("市中区", "市中区"),
        ("资中", "资中县"),
        ("威远", "威远县"),
    ]
    for pat, dist_name in district_patterns:
        if pat in fname:
            district = dist_name
            break
    
    # 从路径中的文件夹名判断区县
    if not district and len(parts) > 2:
        for pat, dist_name in district_patterns:
            if pat in parts[1]:
                district = dist_name
                break
    
    # 市级文件（没有区县）
    if not district:
        # 检查是否是市级文件
        if city == "内江市" and ("内江市委员会" in fname or "neijiang" in fname.lower() or "内江-" in fname):
            district = "市级"
        elif city == "自贡市" and ("自贡-" in fname or "zigong_" in fname.lower()):
            district = "市级"
    
    # 判断单位类型
    if "党校" in fname or "dangxiao" in fname.lower():
        unit_type = "党校"
    elif "组织部" in fname or "zuzhibu" in fname.lower():
        unit_type = "组织部"
    elif "总" in fname and "东部新区" in fname:
        unit_type = "综合"
    
    return city, district, unit_type

def parse_excel(filepath):
    """解析Excel文件，返回所有Sheet的所有行数据"""
    import openpyxl
    results = []
    try:
        wb = openpyxl.load_workbook(filepath, data_only=True)
        for sheet_name in wb.sheetnames:
            ws = wb[sheet_name]
            for row_idx, row in enumerate(ws.iter_rows(values_only=False), 1):
                row_data = []
                for col_idx, cell in enumerate(row, 1):
                    val = cell.value
                    if val is not None:
                        row_data.append({
                            'col': col_idx,
                            'row': row_idx,
                            'value': str(val).strip(),
                            'col_letter': cell.column_letter,
                            'sheet': sheet_name
                        })
                if row_data:
                    results.append({
                        'sheet': sheet_name,
                        'row': row_idx,
                        'cells': row_data
                    })
    except Exception as e:
        # 尝试xlrd
        try:
            import xlrd
            wb = xlrd.open_workbook(filepath)
            for sheet_name in wb.sheet_names():
                ws = wb.sheet_by_name(sheet_name)
                for row_idx in range(ws.nrows):
                    row_data = []
                    for col_idx in range(ws.ncols):
                        val = ws.cell_value(row_idx, col_idx)
                        if val:
                            row_data.append({
                                'col': col_idx + 1,
                                'row': row_idx + 1,
                                'value': str(val).strip(),
                                'col_letter': '',
                                'sheet': sheet_name
                            })
                    if row_data:
                        results.append({
                            'sheet': sheet_name,
                            'row': row_idx + 1,
                            'cells': row_data
                        })
        except Exception as e2:
            print(f"  Excel解析失败 {filepath}: {e} / {e2}")
    return results

def find_amount_in_excel_row(row_data, all_rows, sheet_name):
    """在Excel行数据中查找预算金额"""
    # 策略1: 查找表头中有"预算数"、"总计"、"合计"、"金额"等列
    amount = None
    amount_col = None
    
    # 检查当前行每个单元格
    for cell in row_data:
        val = cell['value']
        # 尝试直接解析数字
        try:
            num = float(val.replace(',', '').replace('，', ''))
            if num > 0 and num < 100000:  # 合理范围
                # 检查是否是金额列（通过表头判断）
                if amount is None:
                    amount = num
                    amount_col = cell['col']
        except (ValueError, TypeError):
            pass
    
    return amount

def find_amount_with_header(row_data, header_row, sheet_name):
    """使用表头行查找金额列"""
    if not header_row:
        return None, None
    
    # 找到金额相关列
    amount_cols = []
    for cell in header_row:
        val = cell['value']
        if any(kw in val for kw in ['预算数', '预算金额', '合计', '总计', '金额', '预算']):
            amount_cols.append(cell['col'])
    
    if not amount_cols:
        return None, None
    
    # 在当前行找到对应列的值
    for cell in row_data:
        if cell['col'] in amount_cols:
            try:
                num = float(cell['value'].replace(',', '').replace('，', ''))
                if num > 0:
                    return num, cell['col']
            except (ValueError, TypeError):
                pass
    
    return None, None

def match_keywords_excel(filepath):
    """解析Excel并匹配关键字"""
    matches = []
    rows_data = parse_excel(filepath)
    
    for row_info in rows_data:
        sheet_name = row_info['sheet']
        row_idx = row_info['row']
        cells = row_info['cells']
        
        # 合并所有单元格文本
        row_text = ' '.join([c['value'] for c in cells])
        
        # 检查是否是表头行（用于金额列识别）
        is_header = any(kw in row_text for kw in ['预算数', '预算金额', '合计', '总计', '金额', '项目', '科目'])
        
        # 查找金额
        amount = None
        # 尝试从当前行找数字
        for cell in cells:
            try:
                num = float(cell['value'].replace(',', '').replace('，', ''))
                if 0 < num < 100000:
                    amount = num
                    break
            except (ValueError, TypeError):
                pass
        
        # 匹配关键字
        for kw in KEYWORDS:
            if kw in row_text:
                # 提取匹配内容（包含关键字的上下文）
                idx = row_text.find(kw)
                start = max(0, idx - 20)
                end = min(len(row_text), idx + len(kw) + 30)
                context = row_text[start:end].strip()
                
                # 构建位置信息
                location = f"{sheet_name} R{row_idx}"
                
                matches.append({
                    'keyword': kw,
                    'content': context,
                    'location': location,
                    'amount': amount if amount else 0.0,
                    'sheet': sheet_name,
                    'row': row_idx
                })
    
    return matches

def parse_docx(filepath):
    """解析docx文件"""
    from docx import Document
    results = []
    try:
        doc = Document(filepath)
        # 段落
        for i, para in enumerate(doc.paragraphs):
            if para.text.strip():
                results.append({
                    'type': 'paragraph',
                    'index': i + 1,
                    'text': para.text.strip()
                })
        # 表格
        for t_idx, table in enumerate(doc.tables):
            for r_idx, row in enumerate(table.rows):
                row_text = ' '.join([cell.text.strip() for cell in row.cells])
                if row_text.strip():
                    results.append({
                        'type': 'table',
                        'table': t_idx + 1,
                        'row': r_idx + 1,
                        'text': row_text.strip()
                    })
    except Exception as e:
        print(f"  docx解析失败 {filepath}: {e}")
    return results

def parse_doc(filepath):
    """解析.doc文件，先转成文本"""
    results = []
    try:
        # 用LibreOffice转换
        with tempfile.TemporaryDirectory() as tmpdir:
            subprocess.run(
                ['libreoffice', '--headless', '--convert-to', 'txt:Text', '--outdir', tmpdir, filepath],
                capture_output=True, timeout=60
            )
            basename = os.path.splitext(os.path.basename(filepath))[0]
            txt_path = os.path.join(tmpdir, basename + '.txt')
            if os.path.exists(txt_path):
                with open(txt_path, 'r', encoding='utf-8', errors='ignore') as f:
                    for i, line in enumerate(f):
                        if line.strip():
                            results.append({
                                'type': 'text',
                                'index': i + 1,
                                'text': line.strip()
                            })
            else:
                print(f"  doc转换失败 {filepath}: 未生成txt")
    except Exception as e:
        print(f"  doc解析失败 {filepath}: {e}")
    return results

def match_keywords_doc(filepath, file_ext):
    """解析Word文件并匹配关键字"""
    if file_ext == '.docx':
        items = parse_docx(filepath)
    else:
        items = parse_doc(filepath)
    
    matches = []
    for item in items:
        text = item['text']
        loc_prefix = f"{item['type']}"
        if item['type'] == 'table':
            loc_prefix = f"表{item['table']} R{item['row']}"
        else:
            loc_prefix = f"R{item['index']}"
        
        # 查找金额
        amount = None
        # 正则匹配 "XX万元" 或 "XX.XX万元"
        amt_matches = re.findall(r'(\d+\.?\d*)\s*万元?', text)
        if amt_matches:
            try:
                amt = float(amt_matches[-1])  # 取最后一个
                if 0 < amt < 100000:
                    amount = amt
            except ValueError:
                pass
        
        # 如果没找到"万元"，尝试找纯数字
        if amount is None:
            nums = re.findall(r'(?<![\d年月日号])(\d+\.?\d*)', text)
            for n in nums:
                try:
                    num = float(n)
                    if 0 < num < 100000 and '.' in n:  # 有小数点的更可能是金额
                        amount = num
                        break
                except ValueError:
                    pass
        
        for kw in KEYWORDS:
            if kw in text:
                idx = text.find(kw)
                start = max(0, idx - 30)
                end = min(len(text), idx + len(kw) + 40)
                context = text[start:end].strip()
                
                matches.append({
                    'keyword': kw,
                    'content': context,
                    'location': loc_prefix,
                    'amount': amount if amount else 0.0,
                    'sheet': '',
                    'row': item.get('index', item.get('row', 0))
                })
    
    return matches

def parse_pdf(filepath):
    """解析PDF文件"""
    results = []
    # 先用pdftotext快速提取
    try:
        with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as tmp:
            tmp_txt = tmp.name
        subprocess.run(
            ['pdftotext', '-layout', filepath, tmp_txt],
            capture_output=True, timeout=60
        )
        if os.path.exists(tmp_txt):
            with open(tmp_txt, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
            lines = content.split('\n')
            for i, line in enumerate(lines):
                if line.strip():
                    results.append({
                        'type': 'text',
                        'page': 1,  # pdftotext不区分页码
                        'line': i + 1,
                        'text': line.strip()
                    })
            os.unlink(tmp_txt)
            
            if not results:
                # 可能是扫描版PDF，用OCR
                results = parse_pdf_ocr(filepath)
    except Exception as e:
        print(f"  PDF解析失败 {filepath}: {e}")
        results = parse_pdf_ocr(filepath)
    
    return results

def parse_pdf_ocr(filepath):
    """OCR解析扫描版PDF"""
    results = []
    try:
        from pdf2image import convert_from_path
        import pytesseract
        
        images = convert_from_path(filepath, dpi=200)
        for page_num, img in enumerate(images, 1):
            text = pytesseract.image_to_string(img, lang='chi_sim')
            lines = text.split('\n')
            for i, line in enumerate(lines):
                if line.strip():
                    results.append({
                        'type': 'ocr',
                        'page': page_num,
                        'line': i + 1,
                        'text': line.strip()
                    })
    except Exception as e:
        print(f"  PDF OCR失败 {filepath}: {e}")
    return results

def match_keywords_pdf(filepath):
    """解析PDF并匹配关键字"""
    items = parse_pdf(filepath)
    matches = []
    
    for item in items:
        text = item['text']
        loc_prefix = f"P{item.get('page', 1)} L{item.get('line', item.get('index', 0))}"
        
        # 查找金额
        amount = None
        amt_matches = re.findall(r'(\d+\.?\d*)\s*万元?', text)
        if amt_matches:
            try:
                amt = float(amt_matches[-1])
                if 0 < amt < 100000:
                    amount = amt
            except ValueError:
                pass
        
        if amount is None:
            nums = re.findall(r'(\d+\.?\d*)', text)
            for n in nums:
                try:
                    num = float(n)
                    if 0 < num < 100000 and '.' in n:
                        amount = num
                        break
                except ValueError:
                    pass
        
        for kw in KEYWORDS:
            if kw in text:
                idx = text.find(kw)
                start = max(0, idx - 30)
                end = min(len(text), idx + len(kw) + 40)
                context = text[start:end].strip()
                
                matches.append({
                    'keyword': kw,
                    'content': context,
                    'location': loc_prefix,
                    'amount': amount if amount else 0.0,
                    'sheet': '',
                    'row': item.get('line', item.get('index', 0))
                })
    
    return matches

def main():
    all_results = {}
    file_count = 0
    total_matches = 0
    
    # 遍历所有文件
    for root, dirs, files in os.walk(BASE_DIR):
        for fname in sorted(files):
            # 跳过脚本自身和已生成的HTML
            if fname.endswith('.py') or fname.endswith('.html') or fname.endswith('.json'):
                continue
            
            filepath = os.path.join(root, fname)
            ext = os.path.splitext(fname)[1].lower()
            
            city, district, unit_type = classify_file(filepath)
            if not city:
                print(f"  ⚠ 无法分类: {filepath}")
                continue
            
            print(f"\n处理: {filepath}")
            print(f"  分类: {city} > {district} > {unit_type}")
            
            matches = []
            if ext in ['.xlsx', '.xls', '.et']:
                matches = match_keywords_excel(filepath)
            elif ext == '.docx':
                matches = match_keywords_doc(filepath, '.docx')
            elif ext == '.doc':
                matches = match_keywords_doc(filepath, '.doc')
            elif ext == '.pdf':
                matches = match_keywords_pdf(filepath)
            else:
                print(f"  跳过未知格式: {ext}")
                continue
            
            print(f"  匹配: {len(matches)} 条")
            total_matches += len(matches)
            file_count += 1
            
            # 构建数据结构
            key = f"{city}||{district}||{unit_type}"
            if key not in all_results:
                all_results[key] = {
                    'city': city,
                    'district': district,
                    'unit_type': unit_type,
                    'files': []
                }
            
            # 相对路径用于文件链接
            rel_path = os.path.relpath(filepath, BASE_DIR)
            
            all_results[key]['files'].append({
                'filename': fname,
                'filepath': filepath,
                'rel_path': rel_path,
                'matches': matches
            })
    
    # 保存为JSON
    output_path = os.path.join(BASE_DIR, "analysis_results.json")
    # 转为可序列化格式
    serializable = {}
    for key, val in all_results.items():
        serializable[key] = val
    
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(serializable, f, ensure_ascii=False, indent=2)
    
    print(f"\n{'='*60}")
    print(f"分析完成!")
    print(f"  文件数: {file_count}")
    print(f"  总匹配: {total_matches}")
    print(f"  分类数: {len(serializable)}")
    print(f"  结果保存: {output_path}")
    
    # 打印每个分类的匹配数
    for key, val in sorted(serializable.items()):
        total = sum(len(f['matches']) for f in val['files'])
        total_amt = sum(m['amount'] for f in val['files'] for m in f['matches'])
        print(f"  {val['city']} > {val['district']} > {val['unit_type']}: {total}条, {total_amt:.2f}万元")

if __name__ == '__main__':
    main()