#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""四川区县预算文件分析 V2 - 修复金额提取
关键修复：
1. 去掉\xa0等不可见字符再匹配
2. 扫描所有列找关键字，不固定name_col
3. 找到关键字后，同行往右找第一个数字作为金额
4. 按市州→区县→单位类型(组织部/党校)分类
5. 同关键字合并，金额累加
"""

import os
import sys
import json
import re
import traceback
from collections import defaultdict

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

CITIES = ['成都','自贡','攀枝花','泸州','德阳','绵阳','广元','遂宁','内江',
          '乐山','南充','眉山','宜宾','广安','达州','巴中','雅安','资阳',
          '阿坝','甘孜','凉山']

DATA_DIR = '/data/www/files/scQxYjs'
OUTPUT_JSON = '/data/www/files/scQxYjs/_final_v2.json'

def clean_str(s):
    """清理字符串：去掉\xa0、多余空格"""
    if s is None:
        return ''
    s = str(s)
    s = s.replace('\xa0', ' ').replace('\u3000', ' ')
    s = re.sub(r'\s+', ' ', s).strip()
    return s

def try_float(val):
    """尝试转为float，失败返回0"""
    if val is None:
        return 0
    if isinstance(val, (int, float)):
        return float(val)
    s = clean_str(val)
    # 去掉逗号
    s = s.replace(',', '').replace('，', '')
    try:
        return float(s)
    except:
        # 提取数字
        m = re.search(r'([\d.]+)', s)
        if m:
            try:
                return float(m.group(1))
            except:
                pass
    return 0

def find_amount_in_row(ws, row_idx, start_col, max_col):
    """在指定行的start_col到max_col之间，找第一个合理的数字作为金额"""
    for c in range(start_col, max_col + 1):
        v = ws.cell(row=row_idx, column=c).value
        if v is not None:
            fv = try_float(v)
            if fv > 0 and fv < 1000000:  # 合理金额范围（万元）
                return fv, c
    return 0, None

# ============ 文件分类 ============
def classify_file(filepath):
    rel = os.path.relpath(filepath, DATA_DIR)
    parts = rel.replace('\\','/').split('/')
    filename = os.path.basename(filepath)
    
    city = ''
    for c in CITIES:
        if c in rel:
            city = c
            break
    
    unit_type = ''
    if '组织部' in filename or '组织部' in rel:
        unit_type = '组织部'
    elif '党校' in filename or '党校' in rel:
        unit_type = '党校'
    else:
        unit_type = '其他'
    
    # 区县
    county = ''
    if len(parts) >= 3:
        # 有子目录
        dir_name = parts[1]
        # 去掉"组织部、党校"等
        county = re.sub(r'组织部.*$|党校.*$', '', dir_name).strip()
        if county.startswith(city):
            county = county[len(city):].strip()
        if county.startswith('市'):
            county = county[1:]
    
    if not county:
        # 从文件名提取
        basename = os.path.splitext(filename)[0]
        name = basename
        if city and name.startswith(city):
            name = name[len(city):]
        # 去掉"市""县""区"前缀
        if name.startswith('市'):
            name = name[1:]
        name = name.replace('组织部','').replace('党校','')
        name = re.sub(r'20\d{2}.*$', '', name)
        name = re.sub(r'预算.*$', '', name)
        name = re.sub(r'部门.*$', '', name)
        name = name.strip('_- ')
        county = name
    
    return city, county, unit_type

# ============ Excel解析 ============
def analyze_excel(filepath):
    matches = []
    try:
        import openpyxl
        wb = openpyxl.load_workbook(filepath, data_only=True)
    except Exception as e:
        print(f'  openpyxl失败: {e}')
        return matches
    
    for sheet_name in wb.sheetnames:
        ws = wb[sheet_name]
        max_r = ws.max_row
        max_c = ws.max_column
        if max_r < 2 or max_c < 2:
            continue
        
        for row_idx in range(1, max_r + 1):
            # 扫描整行所有列，找关键字
            for col_idx in range(1, max_c + 1):
                raw = ws.cell(row=row_idx, column=col_idx).value
                if raw is None:
                    continue
                text = clean_str(raw)
                if not text or len(text) < 2:
                    continue
                
                for kw in KEYWORDS:
                    if kw in text:
                        # 找金额：从当前列往右找第一个数字
                        amount, amt_col = find_amount_in_row(ws, row_idx, col_idx + 1, max_c)
                        # 如果右边没找到，往左找
                        if amount == 0:
                            for c in range(col_idx - 1, 0, -1):
                                v = ws.cell(row=row_idx, column=c).value
                                if v is not None:
                                    fv = try_float(v)
                                    if fv > 0 and fv < 1000000:
                                        amount = fv
                                        amt_col = c
                                        break
                        
                        col_letter = ''
                        try:
                            col_letter = openpyxl.utils.get_column_letter(amt_col) if amt_col else '?'
                        except:
                            col_letter = f'C{amt_col}'
                        
                        matches.append({
                            'keyword': kw,
                            'content': text[:120],
                            'amount': round(amount, 2),
                            'location': f'{sheet_name}:R{row_idx}:{col_letter}',
                            'file': os.path.basename(filepath)
                        })
                        break  # 一个单元格只匹配一个关键字
    
    return matches

def analyze_xls(filepath):
    matches = []
    try:
        import xlrd
        wb = xlrd.open_workbook(filepath)
    except:
        return matches
    
    for sheet in wb.sheets():
        for row_idx in range(sheet.nrows):
            for col_idx in range(min(sheet.ncols, 40)):
                raw = sheet.cell_value(row_idx, col_idx)
                text = clean_str(raw)
                if not text or len(text) < 2:
                    continue
                
                for kw in KEYWORDS:
                    if kw in text:
                        # 找金额：往右找
                        amount = 0
                        amt_col = -1
                        for c in range(col_idx + 1, min(sheet.ncols, 40)):
                            v = sheet.cell_value(row_idx, c)
                            fv = try_float(v)
                            if 0 < fv < 1000000:
                                amount = fv
                                amt_col = c
                                break
                        # 往左找
                        if amount == 0:
                            for c in range(col_idx - 1, -1, -1):
                                v = sheet.cell_value(row_idx, c)
                                fv = try_float(v)
                                if 0 < fv < 1000000:
                                    amount = fv
                                    amt_col = c
                                    break
                        
                        matches.append({
                            'keyword': kw,
                            'content': text[:120],
                            'amount': round(amount, 2),
                            'location': f'{sheet.name}:R{row_idx+1}:C{amt_col+1 if amt_col >= 0 else 0}',
                            'file': os.path.basename(filepath)
                        })
                        break
    
    return matches

# ============ Word解析 ============
def analyze_word(filepath):
    matches = []
    try:
        from docx import Document
        doc = Document(filepath)
    except:
        return analyze_doc_old(filepath)
    
    # 表格
    for table_idx, table in enumerate(doc.tables):
        for row_idx, row in enumerate(table.rows):
            cells = row.cells
            for cell_idx, cell in enumerate(cells):
                text = clean_str(cell.text)
                if not text or len(text) < 2:
                    continue
                for kw in KEYWORDS:
                    if kw in text:
                        # 找金额：往右找
                        amount = 0
                        for c in range(cell_idx + 1, len(cells)):
                            v = try_float(cells[c].text)
                            if 0 < v < 1000000:
                                amount = v
                                break
                        # 往左找
                        if amount == 0:
                            for c in range(cell_idx - 1, -1, -1):
                                v = try_float(cells[c].text)
                                if 0 < v < 1000000:
                                    amount = v
                                    break
                        matches.append({
                            'keyword': kw,
                            'content': text[:120],
                            'amount': round(amount, 2),
                            'location': f'表格{table_idx+1}:行{row_idx+1}',
                            'file': os.path.basename(filepath)
                        })
                        break
    
    # 段落
    for para_idx, para in enumerate(doc.paragraphs):
        text = clean_str(para.text)
        if not text or len(text) < 2:
            continue
        for kw in KEYWORDS:
            if kw in text:
                amount = 0
                m = re.search(r'([\d.]+)\s*万?元', text)
                if m:
                    amount = float(m.group(1))
                matches.append({
                    'keyword': kw,
                    'content': text[:120],
                    'amount': round(amount, 2),
                    'location': f'段落{para_idx+1}',
                    'file': os.path.basename(filepath)
                })
                break
    
    return matches

def analyze_doc_old(filepath):
    matches = []
    try:
        import subprocess
        result = subprocess.run(['antiword', filepath], capture_output=True, text=True, timeout=30)
        text = result.stdout
    except:
        return matches
    
    lines = text.split('\n')
    for line_idx, line in enumerate(lines):
        line = clean_str(line)
        if not line:
            continue
        for kw in KEYWORDS:
            if kw in line:
                amount = 0
                m = re.search(r'([\d.]+)\s*万?元', line)
                if m:
                    amount = float(m.group(1))
                matches.append({
                    'keyword': kw,
                    'content': line[:120],
                    'amount': round(amount, 2),
                    'location': f'行{line_idx+1}',
                    'file': os.path.basename(filepath)
                })
                break
    return matches

# ============ PDF解析 ============
def analyze_pdf(filepath):
    matches = []
    text_pages = []
    
    try:
        import pdfplumber
        with pdfplumber.open(filepath) as pdf:
            for page_idx, page in enumerate(pdf.pages):
                # 先提取表格
                tables = page.extract_tables()
                if tables:
                    for table_idx, table in enumerate(tables):
                        for row_idx, row in enumerate(table):
                            for cell_idx, cell in enumerate(row):
                                text = clean_str(cell)
                                if not text or len(text) < 2:
                                    continue
                                for kw in KEYWORDS:
                                    if kw in text:
                                        # 找金额
                                        amount = 0
                                        for c in range(cell_idx + 1, len(row)):
                                            v = try_float(row[c])
                                            if 0 < v < 1000000:
                                                amount = v
                                                break
                                        if amount == 0:
                                            for c in range(cell_idx - 1, -1, -1):
                                                v = try_float(row[c])
                                                if 0 < v < 1000000:
                                                    amount = v
                                                    break
                                        matches.append({
                                            'keyword': kw,
                                            'content': text[:120],
                                            'amount': round(amount, 2),
                                            'location': f'第{page_idx+1}页:表格{table_idx+1}:行{row_idx+1}',
                                            'file': os.path.basename(filepath)
                                        })
                                        break
                
                # 提取文本
                page_text = page.extract_text()
                if page_text:
                    text_pages.append((page_idx + 1, page_text))
    except Exception as e:
        print(f'  pdfplumber失败: {e}')
    
    # OCR
    if not text_pages and not matches:
        print(f'  尝试OCR: {os.path.basename(filepath)}')
        try:
            from pdf2image import convert_from_path
            import pytesseract
            images = convert_from_path(filepath, dpi=200)
            for page_idx, img in enumerate(images):
                text = pytesseract.image_to_string(img, lang='chi_sim')
                text_pages.append((page_idx + 1, text))
        except Exception as e:
            print(f'  OCR失败: {e}')
    
    # 从文本行匹配
    for page_num, page_text in text_pages:
        lines = page_text.split('\n')
        for line_idx, line in enumerate(lines):
            line = clean_str(line)
            if not line or len(line) < 2:
                continue
            for kw in KEYWORDS:
                if kw in line:
                    amount = 0
                    m = re.search(r'([\d.]+)\s*万?元?', line)
                    if m:
                        try:
                            amount = float(m.group(1))
                        except:
                            pass
                    matches.append({
                        'keyword': kw,
                        'content': line[:120],
                        'amount': round(amount, 2),
                        'location': f'第{page_num}页:行{line_idx+1}',
                        'file': os.path.basename(filepath)
                    })
                    break
    
    # 去重
    seen = set()
    deduped = []
    for m in matches:
        key = (m['keyword'], m['content'][:30])
        if key not in seen:
            seen.add(key)
            deduped.append(m)
    
    return deduped

# ============ ET ============
def analyze_et(filepath):
    try:
        return analyze_excel(filepath)
    except:
        pass
    try:
        return analyze_xls(filepath)
    except:
        pass
    return []

# ============ 主函数 ============
def main():
    # 收集所有文件
    all_files = []
    for root, dirs, files in os.walk(DATA_DIR):
        for f in files:
            if f.startswith('.') or f.startswith('_') or f.endswith('.json') or f.endswith('.html') or f.endswith('.py') or f.endswith('.txt'):
                continue
            filepath = os.path.join(root, f)
            ext = os.path.splitext(f)[1].lower()
            if ext in ['.xlsx', '.xls', '.doc', '.docx', '.pdf', '.et', '.wps']:
                all_files.append(filepath)
    
    print(f'共{len(all_files)}个文件待分析')
    
    # 分析
    results = {}
    processed = 0
    errors = 0
    
    for filepath in all_files:
        processed += 1
        if processed % 50 == 0:
            print(f'进度: {processed}/{len(all_files)}')
        
        ext = os.path.splitext(filepath)[1].lower()
        city, county, unit_type = classify_file(filepath)
        
        if not city:
            continue
        
        if not county:
            county = '未分类'
        
        matches = []
        try:
            if ext == '.xlsx':
                matches = analyze_excel(filepath)
            elif ext == '.xls':
                matches = analyze_xls(filepath)
            elif ext == '.et' or ext == '.wps':
                matches = analyze_et(filepath)
            elif ext == '.docx':
                matches = analyze_word(filepath)
            elif ext == '.doc':
                matches = analyze_doc_old(filepath)
            elif ext == '.pdf':
                matches = analyze_pdf(filepath)
        except Exception as e:
            errors += 1
            if errors <= 10:
                print(f'分析失败: {os.path.basename(filepath)}: {e}')
            continue
        
        if matches:
            rel_path = os.path.relpath(filepath, DATA_DIR)
            for m in matches:
                m['filepath'] = rel_path
            
            if city not in results:
                results[city] = {}
            if county not in results[city]:
                results[city][county] = {'组织部': [], '党校': [], '其他': []}
            if unit_type not in results[city][county]:
                unit_type = '其他'
            results[city][county][unit_type].extend(matches)
    
    print(f'\n分析完成，开始合并关键字...')
    print(f'错误数: {errors}')
    
    # 合并相同关键字
    final = {}
    total_matches = 0
    total_amount = 0
    
    for city, counties in results.items():
        final[city] = {}
        for county, units in counties.items():
            final[city][county] = {}
            for unit_type, items in units.items():
                if not items:
                    continue
                merged = defaultdict(lambda: {'contents': [], 'amounts': [], 'locations': [], 'files': set()})
                for m in items:
                    kw = m['keyword']
                    merged[kw]['contents'].append(m['content'])
                    merged[kw]['amounts'].append(m['amount'])
                    merged[kw]['locations'].append(m['location'])
                    merged[kw]['files'].add(m.get('filepath', m.get('file', '')))
                
                merged_list = []
                unit_total = 0
                for kw, data in merged.items():
                    kw_amount = sum(data['amounts'])
                    unit_total += kw_amount
                    unique_contents = list(dict.fromkeys(data['contents']))
                    merged_list.append({
                        'keyword': kw,
                        'contents': unique_contents,
                        'amount': round(kw_amount, 2),
                        'locations': data['locations'],
                        'files': list(data['files'])
                    })
                    total_matches += 1
                    total_amount += kw_amount
                
                merged_list.sort(key=lambda x: x['amount'], reverse=True)
                final[city][county][unit_type] = {
                    'items': merged_list,
                    'total_amount': round(unit_total, 2),
                    'file_count': len(set().union(*[set(d['files']) for d in [v for v in merged.values()]] if merged else [set()]))
                }
    
    # 保存
    def serialize(obj):
        if isinstance(obj, set):
            return list(obj)
        if isinstance(obj, dict):
            return {k: serialize(v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [serialize(v) for v in obj]
        return obj
    
    with open(OUTPUT_JSON, 'w', encoding='utf-8') as f:
        json.dump(serialize(final), f, ensure_ascii=False, indent=2)
    
    print(f'\n=== 分析完成 ===')
    print(f'总匹配项(合并后): {total_matches}')
    print(f'总金额: {total_amount:.2f}万元')
    
    for city in CITIES:
        if city in final:
            city_amount = 0
            city_items = 0
            for county, units in final[city].items():
                for unit_type, data in units.items():
                    if isinstance(data, dict) and 'items' in data:
                        city_amount += data.get('total_amount', 0)
                        city_items += len(data['items'])
            if city_items > 0:
                print(f'  {city}: {city_items}项, {city_amount:.2f}万元')

if __name__ == '__main__':
    main()