#!/usr/bin/env python3
"""Parse all budget files in 2026-yjs-lzy directory, match 17 keywords, extract amounts."""

import os
import re
import json
import subprocess
import sys
from pathlib import Path

BASE_DIR = "/data/www/files/2026-yjs-lzy"
OUTPUT_JSON = "/data/www/files/2026-yjs-lzy/parsed_all.json"

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

# City/County/Unit type classification
CITY_MAP = {
    "成都市": "成都市", "成都": "成都市",
    "乐山": "乐山市",
    "巴中": "巴中市",
    "攀枝花": "攀枝花市",
    "绵阳": "绵阳市",
    "雅安": "雅安市",
}

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

def classify_file(filepath):
    """Classify file by city, county, unit type based on filename and path."""
    rel_path = os.path.relpath(filepath, BASE_DIR)
    fname = os.path.basename(filepath)
    
    city = None
    county = None
    unit_type = None
    
    # Check if it's in a subdirectory
    parts = rel_path.split("/")
    subdir = parts[0] if len(parts) > 1 else ""
    
    # Determine city from subdir or filename
    for key, city_name in CITY_MAP.items():
        if key in subdir or key in fname:
            city = city_name
            break
    
    # If no city found from direct match, check district names
    if not city:
        for district_key, (city_name, county_name) in DISTRICT_MAP.items():
            if district_key in fname or district_key in subdir:
                city = city_name
                county = county_name
                break
    
    # Determine county
    if not county:
        for district_key, (city_name, county_name) in DISTRICT_MAP.items():
            if district_key in fname or district_key in subdir:
                city = city_name
                county = county_name
                break
    
    # Determine unit type
    if "党校" in fname or "dx_" in fname or "_dx_" in fname:
        unit_type = "党校"
    elif "组织部" in fname or "org_" in fname or "_org_" in fname:
        unit_type = "组织部"
    elif "经济科技" in fname:
        unit_type = "经济科技和信息化局"
    else:
        # Default: check for 党校 or 组织部 in filename
        if "党校" in fname:
            unit_type = "党校"
        elif "组织部" in fname:
            unit_type = "组织部"
        else:
            unit_type = "其他"
    
    # Special: if no county but city is determined, use "市级"
    if not county and city:
        county = "市级"
    
    # If still no city, try harder
    if not city:
        # Check all district keys
        for district_key, (city_name, county_name) in DISTRICT_MAP.items():
            if district_key in fname:
                city = city_name
                county = county_name
                break
        if not city:
            city = "未分类"
            county = "未分类"
    
    return city, county, unit_type

def extract_text_pdf(filepath):
    """Extract text from PDF using pdftotext."""
    try:
        result = subprocess.run(["pdftotext", "-layout", filepath, "-"], 
                              capture_output=True, text=True, timeout=60)
        if result.returncode == 0 and result.stdout.strip():
            return result.stdout
        return ""
    except Exception as e:
        print(f"  [PDF ERROR] {filepath}: {e}", file=sys.stderr)
        return ""

def extract_text_excel(filepath):
    """Extract text from Excel using openpyxl."""
    try:
        from openpyxl import load_workbook
        wb = load_workbook(filepath, data_only=True, read_only=True)
        lines = []
        for sheet in wb.sheetnames:
            ws = wb[sheet]
            lines.append(f"=== Sheet: {sheet} ===")
            for row in ws.iter_rows(values_only=True):
                row_text = []
                for cell in row:
                    if cell is not None:
                        row_text.append(str(cell))
                if row_text:
                    lines.append(" | ".join(row_text))
        wb.close()
        return "\n".join(lines)
    except Exception as e:
        print(f"  [EXCEL ERROR] {filepath}: {e}", file=sys.stderr)
        return ""

def extract_text_word(filepath):
    """Extract text from Word doc."""
    try:
        # Try python-docx first
        if filepath.endswith('.docx'):
            from docx import Document
            doc = Document(filepath)
            lines = []
            for para in doc.paragraphs:
                if para.text.strip():
                    lines.append(para.text)
            for table in doc.tables:
                for row in table.rows:
                    row_text = []
                    for cell in row.cells:
                        if cell.text.strip():
                            row_text.append(cell.text.strip())
                    if row_text:
                        lines.append(" | ".join(row_text))
            return "\n".join(lines)
        else:
            # .doc - use libreoffice
            result = subprocess.run(
                ["libreoffice", "--headless", "--convert-to", "txt:Text", 
                 "--outdir", "/tmp", filepath],
                capture_output=True, text=True, timeout=60
            )
            if result.returncode == 0:
                txt_path = "/tmp/" + os.path.splitext(os.path.basename(filepath))[0] + ".txt"
                if os.path.exists(txt_path):
                    with open(txt_path, 'r', errors='ignore') as f:
                        return f.read()
            return ""
    except Exception as e:
        print(f"  [WORD ERROR] {filepath}: {e}", file=sys.stderr)
        return ""

def extract_text_et(filepath):
    """Extract text from .et file (WPS spreadsheet) - treat as xlsx."""
    return extract_text_excel(filepath)

def extract_amount_from_line(line):
    """Extract budget amount from a line of text."""
    # Pattern: number followed by 万元 or just number
    patterns = [
        r'(\d+\.?\d*)\s*万元',
        r'(\d+\.?\d*)\s*万',
        r'(\d+\.?\d*)\s*元',
    ]
    amounts = []
    for pattern in patterns:
        matches = re.findall(pattern, line)
        for m in matches:
            try:
                amt = float(m)
                if amt > 0:
                    # Convert to 万元
                    if '万元' in line or '万' in line:
                        amounts.append(round(amt, 2))
                    elif '元' in line:
                        amounts.append(round(amt / 10000, 2))
                    else:
                        # Bare number, likely in 元
                        if amt > 100:
                            amounts.append(round(amt / 10000, 2))
                        else:
                            amounts.append(round(amt, 2))
            except ValueError:
                continue
    
    # Also try to find bare numbers in table cells (for Excel)
    # Look for patterns like "123.45" or "12345" that could be budget amounts
    bare_numbers = re.findall(r'(?<!\d)(\d+\.?\d{0,2})(?!\d)', line)
    for num_str in bare_numbers:
        try:
            num = float(num_str)
            if 0.1 <= num <= 50000:  # Reasonable budget range in 万元
                # Only add if no other amount found on this line
                if not amounts:
                    amounts.append(round(num, 2))
        except ValueError:
            continue
    
    # Return the most likely amount (usually the last one or the one closest to 万元)
    if amounts:
        return max(amounts)  # Return the largest as most likely the budget amount
    
    # For Excel table cells, try to find numbers in the line
    numbers = re.findall(r'(\d+\.?\d*)', line)
    for n in numbers:
        try:
            num = float(n)
            if 0.01 <= num <= 50000:
                return round(num, 2)
        except ValueError:
            continue
    
    return None

def match_keywords(text, filepath, city, county, unit_type):
    """Match 17 keywords in text and return list of matches."""
    matches = []
    if not text:
        return matches
    
    lines = text.split('\n')
    
    for keyword in KEYWORDS:
        for line_num, line in enumerate(lines, 1):
            if keyword.lower() in line.lower():
                # Extract amount
                amount = extract_amount_from_line(line)
                
                # Determine location
                ext = os.path.splitext(filepath)[1].lower()
                if ext in ('.xlsx', '.xls', '.et'):
                    # Find sheet name
                    sheet_name = "Sheet1"
                    for i in range(line_num - 1, max(0, line_num - 20), -1):
                        if i < len(lines) and lines[i].startswith("=== Sheet:"):
                            sheet_name = lines[i].replace("=== Sheet:", "").replace("===", "").strip()
                            break
                    location = f"{sheet_name}:R{line_num}"
                elif ext == '.pdf':
                    location = f"P{line_num}"
                else:
                    location = f"L{line_num}"
                
                # Truncate line content for display
                content = line.strip()[:200]
                if not content:
                    continue
                
                matches.append({
                    "keyword": keyword,
                    "content": content,
                    "location": location,
                    "amount": amount if amount else 0.0,
                    "file": os.path.basename(filepath),
                    "filepath": filepath,
                    "city": city,
                    "county": county,
                    "unit_type": unit_type,
                    "line_number": line_num
                })
    
    return matches

def main():
    print("=" * 60)
    print("开始解析 2026-yjs-lzy 目录预算文件")
    print("=" * 60)
    
    # Find all budget files
    extensions = ('.pdf', '.xlsx', '.xls', '.doc', '.docx', '.et')
    files = []
    for root, dirs, filenames in os.walk(BASE_DIR):
        for fname in filenames:
            if fname.lower().endswith(extensions):
                files.append(os.path.join(root, fname))
    
    files.sort()
    print(f"找到 {len(files)} 个预算文件")
    
    all_matches = []
    file_stats = []
    
    for i, filepath in enumerate(files, 1):
        fname = os.path.basename(filepath)
        rel_path = os.path.relpath(filepath, BASE_DIR)
        ext = os.path.splitext(filepath)[1].lower()
        
        print(f"\n[{i}/{len(files)}] 解析: {rel_path}")
        
        # Classify file
        city, county, unit_type = classify_file(filepath)
        print(f"  分类: {city} > {county} > {unit_type}")
        
        # Extract text
        text = ""
        if ext == '.pdf':
            text = extract_text_pdf(filepath)
        elif ext in ('.xlsx', '.xls'):
            text = extract_text_excel(filepath)
        elif ext in ('.doc', '.docx'):
            text = extract_text_word(filepath)
        elif ext == '.et':
            text = extract_text_et(filepath)
        
        if not text:
            print(f"  ⚠ 无法提取文本")
            file_stats.append({"file": fname, "path": rel_path, "matches": 0, "text_length": 0, "city": city, "county": county, "unit_type": unit_type})
            continue
        
        print(f"  文本长度: {len(text)} 字符")
        
        # Match keywords
        matches = match_keywords(text, filepath, city, county, unit_type)
        print(f"  匹配: {len(matches)} 条")
        
        all_matches.extend(matches)
        file_stats.append({
            "file": fname, "path": rel_path, "matches": len(matches),
            "text_length": len(text), "city": city, "county": county, "unit_type": unit_type
        })
    
    # Summary
    print("\n" + "=" * 60)
    print("解析完成!")
    print(f"总文件数: {len(files)}")
    print(f"总匹配数: {len(all_matches)}")
    
    # Keyword distribution
    kw_dist = {}
    for m in all_matches:
        kw = m["keyword"]
        if kw not in kw_dist:
            kw_dist[kw] = {"count": 0, "amount": 0.0}
        kw_dist[kw]["count"] += 1
        kw_dist[kw]["amount"] += m["amount"]
    
    print("\n关键字分布:")
    for kw in KEYWORDS:
        if kw in kw_dist:
            print(f"  {kw}: {kw_dist[kw]['count']}次, {kw_dist[kw]['amount']:.2f}万元")
    
    # City distribution
    city_dist = {}
    for m in all_matches:
        c = m["city"]
        if c not in city_dist:
            city_dist[c] = {"count": 0, "amount": 0.0, "counties": set()}
        city_dist[c]["count"] += 1
        city_dist[c]["amount"] += m["amount"]
        city_dist[c]["counties"].add(m["county"])
    
    print("\n市州分布:")
    for c, d in sorted(city_dist.items()):
        print(f"  {c}: {d['count']}次, {d['amount']:.2f}万元, {len(d['counties'])}个区县")
    
    # Save results
    # Convert sets to lists for JSON
    for c in city_dist:
        city_dist[c]["counties"] = list(city_dist[c]["counties"])
    
    result = {
        "total_files": len(files),
        "total_matches": len(all_matches),
        "total_amount": sum(m["amount"] for m in all_matches),
        "keyword_distribution": kw_dist,
        "city_distribution": city_dist,
        "matches": all_matches,
        "file_stats": file_stats
    }
    
    with open(OUTPUT_JSON, 'w', encoding='utf-8') as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    
    print(f"\n结果已保存: {OUTPUT_JSON}")
    print(f"JSON大小: {os.path.getsize(OUTPUT_JSON) / 1024:.1f} KB")

if __name__ == "__main__":
    main()