#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
四川区县预算文件完整分析脚本
规则：
1. 17个关键字匹配
2. 每条匹配项提取对应预算金额
3. 按市州→区县→单位类型(组织部/党校)分类
4. 同一区县同一单位类型下，相同关键字合并为一行，金额累加
5. 每个区县组织部/党校分别汇总总金额
"""

import os
import re
import json
import traceback
import subprocess
import sys

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

BASE_DIR = '/data/www/files/scQxYjs'
OUTPUT_FILE = os.path.join(BASE_DIR, 'final_results.json')

# 排除的Sheet名（封面、绩效目标表等非预算数据）
EXCLUDE_SHEETS = ['封面', '绩效', '目标']

def should_process_sheet(sheet_name):
    """判断是否需要处理该Sheet"""
    for ex in EXCLUDE_SHEETS:
        if ex in sheet_name:
            return False
    return True

def classify_unit(filename, filepath):
    """根据文件名判断单位类型：组织部/党校"""
    fn = filename.lower()
    base_fn = os.path.basename(filename)
    
    if '党校' in base_fn or '觉校' in base_fn:  # 注意OCR可能把"党校"识别成"觉校"
        return '党校'
    if '组织' in base_fn:
        return '组织部'
    
    # 从路径中找
    path_parts = filepath.replace(BASE_DIR, '').split('/')
    for part in path_parts:
        if '党校' in part or '觉校' in part:
            return '党校'
        if '组织' in part:
            return '组织部'
    
    return '其他'

def extract_county(filepath, filename):
    """从路径中提取区县名"""
    rel_path = os.path.relpath(filepath, BASE_DIR)
    parts = rel_path.split('/')
    
    # 找市州下面的区县目录
    if len(parts) >= 2:
        county = parts[1]  # 如 "广元市朝天区"
        # 清理
        county = county.replace('组织部、党校', '').strip()
        return county
    
    return os.path.splitext(filename)[0]

def extract_city(filepath):
    """从路径中提取市州名"""
    rel_path = os.path.relpath(filepath, BASE_DIR)
    parts = rel_path.split('/')
    if len(parts) >= 1:
        return parts[0]
    return '未知'

def try_float(val):
    """安全转换为float"""
    if val is None:
        return 0.0
    if isinstance(val, (int, float)):
        return float(val)
    s = str(val).strip().replace(' ', '').replace('\xa0', '')
    if not s:
        return 0.0
    # 去掉非数字字符
    s = re.sub(r'[^\d.]', '', s)
    if not s:
        return 0.0
    try:
        return float(s)
    except:
        return 0.0

def find_amount_column(ws, header_row, name_col):
    """根据表头找到金额列（预算数/合计/总计/金额列）"""
    amount_cols = []
    
    # 扫描表头行
    for c in range(1, ws.max_column + 1):
        val = ws.cell(header_row, c).value
        if val is None:
            continue
        val_str = str(val).strip().replace(' ', '').replace('\xa0', '')
        if any(kw in val_str for kw in ['预算数', '合计', '总计', '金额', '预算']):
            if c != name_col:  # 金额列不能和名称列相同
                amount_cols.append(c)
    
    return amount_cols

def find_name_column(ws, header_rows=range(1, 8)):
    """找到项目名称列"""
    for r in header_rows:
        for c in range(1, ws.max_column + 1):
            val = ws.cell(r, c).value
            if val is None:
                continue
            val_str = str(val).strip().replace(' ', '').replace('\xa0', '')
            if any(kw in val_str for kw in ['单位名称', '项目名称', '科目', '名称（科目）', '支出功能分类']):
                return c, r
    return None, None

def analyze_excel(filepath):
    """分析Excel文件"""
    results = []
    try:
        import openpyxl
        wb = openpyxl.load_workbook(filepath, data_only=True)
        
        for sheet_name in wb.sheetnames:
            if not should_process_sheet(sheet_name):
                continue
            
            ws = wb[sheet_name]
            if ws.max_row < 3 or ws.max_column < 2:
                continue
            
            # 找名称列和金额列
            name_col, header_row = find_name_column(ws)
            if name_col is None:
                # 退而求其次：逐行逐列扫描
                for r in range(1, ws.max_row + 1):
                    for c in range(1, ws.max_column + 1):
                        val = ws.cell(r, c).value
                        if val is None:
                            continue
                        val_str = str(val).strip().replace('\xa0', '')
                        if not val_str or val_str == 'None':
                            continue
                        
                        # 匹配关键字
                        for kw in KEYWORDS:
                            if kw in val_str:
                                # 找同行其他列的数字作为金额
                                amount = 0.0
                                for ac in range(1, ws.max_column + 1):
                                    if ac == c:
                                        continue
                                    av = ws.cell(r, ac).value
                                    if av is not None and try_float(av) > 0:
                                        amount = try_float(av)
                                        break
                                
                                results.append({
                                    'keyword': kw,
                                    'content': val_str,
                                    'sheet': sheet_name,
                                    'row': r,
                                    'col': c,
                                    'amount': round(amount, 2)
                                })
                                break  # 一个单元格只匹配第一个关键字
                continue
            
            # 找金额列
            amount_cols = find_amount_column(ws, header_row, name_col)
            
            # 逐行处理
            for r in range(header_row + 1, ws.max_row + 1):
                name_val = ws.cell(r, name_col).value
                if name_val is None:
                    continue
                name_str = str(name_val).strip().replace('\xa0', '')
                if not name_str or name_str == 'None':
                    continue
                
                # 匹配关键字
                for kw in KEYWORDS:
                    if kw in name_str:
                        # 提取金额：优先从金额列，否则从同行其他数字列
                        amount = 0.0
                        if amount_cols:
                            for ac in amount_cols:
                                av = ws.cell(r, ac).value
                                if av is not None:
                                    fv = try_float(av)
                                    if fv > 0:
                                        amount = fv
                                        break
                        
                        if amount == 0.0:
                            # 退而求其次：同行找数字
                            for ac in range(1, ws.max_column + 1):
                                if ac == name_col:
                                    continue
                                av = ws.cell(r, ac).value
                                if av is not None:
                                    fv = try_float(av)
                                    if fv > 0:
                                        amount = fv
                                        break
                        
                        col_letter = ''
                        # 列名
                        col_letter = openpyxl.utils.get_column_letter(name_col) if name_col else ''
                        amount_letter = openpyxl.utils.get_column_letter(amount_cols[0]) if amount_cols else ''
                        
                        results.append({
                            'keyword': kw,
                            'content': name_str,
                            'sheet': sheet_name,
                            'row': r,
                            'col': f'{col_letter}列',
                            'amount_col': f'{amount_letter}列' if amount_cols else '',
                            'amount': round(amount, 2)
                        })
                        break  # 一个项目只匹配第一个关键字
            
            wb.close()
    except Exception as e:
        print(f"  [ERROR] Excel分析失败: {e}")
    
    return results

def analyze_pdf(filepath):
    """分析PDF文件"""
    results = []
    try:
        import pdfplumber
        
        with pdfplumber.open(filepath) as pdf:
            for page_num, page in enumerate(pdf.pages, 1):
                # 先尝试提取表格
                tables = page.extract_tables()
                if tables:
                    for table in tables:
                        for row_idx, row in enumerate(table):
                            if not row:
                                continue
                            for cell_idx, cell in enumerate(row):
                                if cell is None:
                                    continue
                                cell_str = str(cell).strip()
                                if not cell_str:
                                    continue
                                
                                for kw in KEYWORDS:
                                    if kw in cell_str:
                                        # 找同行其他单元格的数字
                                        amount = 0.0
                                        for other_idx, other_cell in enumerate(row):
                                            if other_idx == cell_idx or other_cell is None:
                                                continue
                                            fv = try_float(other_cell)
                                            if fv > 0:
                                                amount = fv
                                                break
                                        
                                        results.append({
                                            'keyword': kw,
                                            'content': cell_str,
                                            'sheet': f'第{page_num}页',
                                            'row': row_idx + 1,
                                            'col': f'第{cell_idx+1}列',
                                            'amount': round(amount, 2)
                                        })
                                        break
                else:
                    # 没有表格，提取文本
                    text = page.extract_text()
                    if not text:
                        continue
                    
                    lines = text.split('\n')
                    for line_idx, line in enumerate(lines):
                        line = line.strip()
                        if not line:
                            continue
                        
                        for kw in KEYWORDS:
                            if kw in line:
                                # 从行中提取金额
                                amount = 0.0
                                # 模式1：数字后跟万元
                                m = re.search(r'(\d+\.?\d*)\s*万?元', line)
                                if m:
                                    amount = float(m.group(1))
                                else:
                                    # 模式2：行中任意数字
                                    nums = re.findall(r'(\d+\.?\d*)', line)
                                    for n in nums:
                                        fv = float(n)
                                        if 0 < fv < 100000:  # 合理范围
                                            amount = fv
                                            break
                                
                                results.append({
                                    'keyword': kw,
                                    'content': line[:100],  # 截取前100字符
                                    'sheet': f'第{page_num}页',
                                    'row': line_idx + 1,
                                    'col': '文本',
                                    'amount': round(amount, 2)
                                })
                                break
    except Exception as e:
        print(f"  [ERROR] PDF分析失败: {e}")
    
    return results

def analyze_doc(filepath):
    """分析Word文件"""
    results = []
    try:
        from docx import Document
        doc = Document(filepath)
        
        # 处理表格
        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]
                for cell_idx, cell_text in enumerate(cells):
                    if not cell_text:
                        continue
                    
                    for kw in KEYWORDS:
                        if kw in cell_text:
                            # 找同行其他单元格的数字
                            amount = 0.0
                            for other_idx, other_text in enumerate(cells):
                                if other_idx == cell_idx:
                                    continue
                                fv = try_float(other_text)
                                if fv > 0:
                                    amount = fv
                                    break
                            
                            results.append({
                                'keyword': kw,
                                'content': cell_text[:100],
                                'sheet': f'表格{table_idx+1}',
                                'row': row_idx + 1,
                                'col': f'第{cell_idx+1}列',
                                'amount': round(amount, 2)
                            })
                            break
        
        # 处理段落文本
        for para_idx, para in enumerate(doc.paragraphs):
            text = para.text.strip()
            if not text:
                continue
            
            for kw in KEYWORDS:
                if kw in text:
                    # 从文本中提取金额
                    amount = 0.0
                    m = re.search(r'(\d+\.?\d*)\s*万?元', text)
                    if m:
                        amount = float(m.group(1))
                    
                    results.append({
                        'keyword': kw,
                        'content': text[:100],
                        'sheet': '段落',
                        'row': para_idx + 1,
                        'col': '文本',
                        'amount': round(amount, 2)
                    })
                    break
    except Exception as e:
        print(f"  [ERROR] Word分析失败: {e}")
    
    return results

def scan_pdf_ocr(filepath):
    """扫描版PDF走OCR"""
    results = []
    try:
        from pdf2image import convert_from_path
        import pytesseract
        from PIL import Image
        
        images = convert_from_path(filepath, dpi=200)
        
        for page_num, image in enumerate(images, 1):
            text = pytesseract.image_to_string(image, lang='chi_sim')
            lines = text.split('\n')
            
            for line_idx, line in enumerate(lines):
                line = line.strip()
                if not line or len(line) < 3:
                    continue
                
                for kw in KEYWORDS:
                    if kw in line:
                        amount = 0.0
                        m = re.search(r'(\d+\.?\d*)\s*万?元', line)
                        if m:
                            amount = float(m.group(1))
                        
                        results.append({
                            'keyword': kw,
                            'content': line[:100],
                            'sheet': f'第{page_num}页(OCR)',
                            'row': line_idx + 1,
                            'col': '文本',
                            'amount': round(amount, 2)
                        })
                        break
    except Exception as e:
        print(f"  [ERROR] OCR分析失败: {e}")
    
    return results

def analyze_file(filepath):
    """分析单个文件"""
    filename = os.path.basename(filepath)
    ext = os.path.splitext(filename)[1].lower()
    
    results = []
    
    if ext in ['.xlsx', '.xls']:
        results = analyze_excel(filepath)
    elif ext == '.et':
        # WPS格式，先尝试openpyxl
        try:
            results = analyze_excel(filepath)
        except:
            results = []
    elif ext == '.pdf':
        results = analyze_pdf(filepath)
        # 如果pdfplumber没提取到内容，走OCR
        if not results:
            print(f"  [INFO] PDF无文本，尝试OCR: {filename}")
            results = scan_pdf_ocr(filepath)
    elif ext in ['.doc', '.docx']:
        results = analyze_doc(filepath)
    else:
        return []
    
    return results

def main():
    # 收集所有文件
    all_files = []
    for root, dirs, files in os.walk(BASE_DIR):
        for f in files:
            if f.startswith('.') or f.startswith('_') or f.startswith('full_') or f.startswith('batch') or f == 'final_results.json':
                continue
            ext = os.path.splitext(f)[1].lower()
            if ext in ['.xlsx', '.xls', '.et', '.pdf', '.doc', '.docx']:
                all_files.append(os.path.join(root, f))
    
    print(f"共找到 {len(all_files)} 个预算文件")
    
    # 分析结果：市州 → 单位类型 → 区县 → 匹配项列表
    data = {}
    total_matches = 0
    
    for i, filepath in enumerate(all_files):
        filename = os.path.basename(filepath)
        city = extract_city(filepath)
        county = extract_county(filepath, filename)
        unit_type = classify_unit(filename, filepath)
        
        if unit_type == '其他':
            continue  # 跳过非组织部/党校文件
        
        print(f"[{i+1}/{len(all_files)}] {city}/{county}/{unit_type}: {filename}")
        
        matches = analyze_file(filepath)
        
        if not matches:
            continue
        
        # 按关键字合并
        merged = {}  # keyword → {content_list, amount_sum, locations}
        for m in matches:
            kw = m['keyword']
            if kw not in merged:
                merged[kw] = {
                    'keyword': kw,
                    'contents': [],
                    'amount': 0.0,
                    'locations': []
                }
            merged[kw]['contents'].append(m['content'])
            merged[kw]['amount'] += m['amount']
            loc = f"{m['sheet']}/R{m['row']}/{m['col']}"
            if m.get('amount_col'):
                loc += f"/金额列:{m['amount_col']}"
            merged[kw]['locations'].append(loc)
        
        # 转为列表
        merged_list = list(merged.values())
        for item in merged_list:
            item['amount'] = round(item['amount'], 2)
            item['contents'] = ' + '.join(item['contents'][:5])  # 最多显示5个
            item['locations'] = ' | '.join(item['locations'][:5])
        
        # 汇总金额
        total_amount = round(sum(item['amount'] for item in merged_list), 2)
        
        # 存入数据结构
        if city not in data:
            data[city] = {}
        if unit_type not in data[city]:
            data[city][unit_type] = {}
        if county not in data[city][unit_type]:
            data[city][unit_type][county] = {}
        
        data[city][unit_type][county] = {
            'filename': filename,
            'filepath': filepath,
            'total_amount': total_amount,
            'match_count': len(merged_list),
            'matches': merged_list
        }
        
        total_matches += len(merged_list)
        
        if (i+1) % 50 == 0:
            print(f"\n--- 已处理 {i+1}/{len(all_files)}，累计匹配 {total_matches} 项 ---\n")
    
    # 保存结果
    with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    
    print(f"\n{'='*60}")
    print(f"分析完成！")
    print(f"总文件数: {len(all_files)}")
    print(f"总匹配项: {total_matches}")
    print(f"结果文件: {OUTPUT_FILE}")
    
    # 按市州统计
    print(f"\n各市州匹配情况:")
    for city in sorted(data.keys()):
        org_count = len(data[city].get('组织部', {}))
        dx_count = len(data[city].get('党校', {}))
        org_amt = sum(data[city]['组织部'][c]['total_amount'] for c in data[city].get('组织部', {}))
        dx_amt = sum(data[city]['党校'][c]['total_amount'] for c in data[city].get('党校', {}))
        print(f"  {city}: 组织部{org_count}个区县/{org_amt:.2f}万, 党校{dx_count}个区县/{dx_amt:.2f}万")

if __name__ == '__main__':
    main()