#!/usr/bin/env python3
"""提取所有PDF预算文件文本，匹配17个关键字，提取预算金额"""
import os
import re
import json
import subprocess
import sys

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

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

# 文件分类映射（根据文件名判断市州/区县/单位类型）
FILE_CLASSIFICATION = {
    "西宁市2025年预算执行情况和2026年预算.pdf": ("西宁市", "市级", "政府"),
    "海东市-党校-2.pdf": ("海东市", "市级", "党校"),
    "西宁市2025年预算执行情况和2026年预算草案的报告.pdf": ("西宁市", "市级", "政府"),
    "海东市-组织部-1.pdf": ("海东市", "市级", "组织部"),
    "xining_zuzhibu_2026.pdf": ("西宁市", "市级", "组织部"),
    "海东市-组织部-2.pdf": ("海东市", "市级", "组织部"),
    "西宁市民政局2026年单位预算.pdf": ("西宁市", "市级", "民政局"),
    "西宁市城中区民政局2026预算.pdf": ("西宁市", "城中区", "民政局"),
    "西宁市民政局2026年部门预算.pdf": ("西宁市", "市级", "民政局"),
    "中共青海省委组织部2026年部门预算.pdf": ("青海省", "省级", "组织部"),
    "西宁市城东区民政局2026年部门预算.pdf": ("西宁市", "城东区", "民政局"),
    "xining_dangxiao_2026.pdf": ("西宁市", "市级", "党校"),
    "海东市-党校-1.pdf": ("海东市", "市级", "党校"),
    "2026年西宁市政府预算公开目录.pdf": ("西宁市", "市级", "政府"),
}

def extract_pdf_text(filepath):
    """用pdftotext提取PDF文本"""
    try:
        result = subprocess.run(
            ["pdftotext", "-layout", filepath, "-"],
            capture_output=True, text=True, timeout=120
        )
        return result.stdout
    except Exception as e:
        print(f"  [ERROR] pdftotext failed: {e}", file=sys.stderr)
        return ""

def extract_pdf_pages(filepath):
    """用pdftotext逐页提取PDF文本"""
    try:
        # 先获取页数
        result = subprocess.run(
            ["pdfinfo", filepath],
            capture_output=True, text=True, timeout=30
        )
        pages = 1
        for line in result.stdout.split('\n'):
            if line.startswith('Pages:'):
                pages = int(line.split(':')[1].strip())
                break
    except:
        pages = 1
    
    all_pages = []
    for p in range(1, pages + 1):
        try:
            result = subprocess.run(
                ["pdftotext", "-layout", "-f", str(p), "-l", str(p), filepath, "-"],
                capture_output=True, text=True, timeout=60
            )
            all_pages.append((p, result.stdout))
        except:
            all_pages.append((p, ""))
    return all_pages

def match_keywords_in_text(text, page_num=None):
    """在文本中匹配17个关键字，提取预算金额"""
    lines = text.split('\n')
    matches = []
    
    for line_idx, line in enumerate(lines):
        line_stripped = line.strip()
        if not line_stripped:
            continue
        
        for kw in KEYWORDS:
            if kw in line_stripped:
                # 提取金额
                amount = extract_amount_from_line(line_stripped, lines, line_idx)
                
                # 构建位置信息
                if page_num:
                    location = f"P{page_num}:L{line_idx+1}"
                else:
                    location = f"L{line_idx+1}"
                
                matches.append({
                    "keyword": kw,
                    "content": line_stripped[:200],
                    "location": location,
                    "amount": amount,
                    "page": page_num or 1
                })
    
    return matches

def extract_amount_from_line(line, all_lines, line_idx):
    """从行中提取预算金额"""
    # 策略1: 行内直接有"XX万元"或"XX.XX万元"
    patterns = [
        r'(\d+\.?\d*)\s*万元',
        r'(\d+\.?\d*)\s*万',
    ]
    for pat in patterns:
        m = re.search(pat, line)
        if m:
            val = float(m.group(1))
            if 0.01 <= val <= 10000:
                return val
    
    # 策略2: 行内有数字，可能是预算数列
    # 找行内所有数字
    numbers = re.findall(r'(\d+\.?\d*)', line)
    for num_str in numbers:
        val = float(num_str)
        # 排除年份、百分比、明显非金额的数字
        if val > 2027 or val < 0.01:
            continue
        if val == 2026 or val == 2025 or val == 2024:
            continue
        if 0.01 <= val <= 10000:
            return val
    
    # 策略3: 看下一行是否有金额
    if line_idx + 1 < len(all_lines):
        next_line = all_lines[line_idx + 1].strip()
        for pat in patterns:
            m = re.search(pat, next_line)
            if m:
                val = float(m.group(1))
                if 0.01 <= val <= 10000:
                    return val
    
    return 0.0

def main():
    files = sorted([f for f in os.listdir(BASE_DIR) if f.endswith('.pdf')])
    print(f"Found {len(files)} PDF files")
    
    all_results = {}
    
    for fname in files:
        fpath = os.path.join(BASE_DIR, fname)
        classification = FILE_CLASSIFICATION.get(fname, ("未分类", "未知", "未知"))
        city, district, org_type = classification
        
        print(f"\n{'='*60}")
        print(f"Processing: {fname}")
        print(f"  City: {city}, District: {district}, OrgType: {org_type}")
        
        # 逐页提取
        pages = extract_pdf_pages(fpath)
        print(f"  Pages: {len(pages)}")
        
        all_matches = []
        for page_num, page_text in pages:
            if page_text.strip():
                matches = match_keywords_in_text(page_text, page_num)
                all_matches.extend(matches)
        
        # 去重（相同关键字+相同行）
        seen = set()
        unique_matches = []
        for m in all_matches:
            key = (m["keyword"], m["content"][:50], m["location"])
            if key not in seen:
                seen.add(key)
                unique_matches.append(m)
        
        print(f"  Total matches: {len(unique_matches)}")
        for m in unique_matches[:5]:
            print(f"    [{m['keyword']}] {m['content'][:60]}... → {m['amount']}万元")
        if len(unique_matches) > 5:
            print(f"    ... and {len(unique_matches)-5} more")
        
        all_results[fname] = {
            "city": city,
            "district": district,
            "org_type": org_type,
            "pages": len(pages),
            "matches": unique_matches
        }
    
    # 保存JSON
    output_path = "/data/www/files/2026-yjs-dsl/analysis_results.json"
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(all_results, f, ensure_ascii=False, indent=2)
    
    print(f"\n{'='*60}")
    print(f"Results saved to: {output_path}")
    
    # 汇总统计
    total_matches = sum(len(r["matches"]) for r in all_results.values())
    total_amount = sum(m["amount"] for r in all_results.values() for m in r["matches"])
    cities = set(r["city"] for r in all_results.values())
    districts = set((r["city"], r["district"]) for r in all_results.values())
    org_types = set(r["org_type"] for r in all_results.values())
    kw_used = set(m["keyword"] for r in all_results.values() for m in r["matches"])
    
    print(f"\nSummary:")
    print(f"  Files: {len(all_results)}")
    print(f"  Total matches: {total_matches}")
    print(f"  Total amount: {total_amount:.2f}万元")
    print(f"  Cities: {len(cities)} - {cities}")
    print(f"  Districts: {len(districts)}")
    print(f"  Org types: {org_types}")
    print(f"  Keywords used: {len(kw_used)} - {kw_used}")

if __name__ == "__main__":
    main()