# figure_table_pipeline.py
import os
import shutil
import re
from pathlib import Path
from collections import defaultdict
from pragent.backend.loader import ImagePDFLoader
from pragent.backend.yolo import extract_and_save_layout_components
from tqdm.asyncio import tqdm

import asyncio
from typing import Any

async def run_figure_extraction(pdf_path: str, base_work_dir: str, progress: Any | None = None) -> str:
    """
    一个完整的、从PDF提取并配对图表的流程。
    这是被 app.py 调用的主函数。

    Args:
        pdf_path (str): 用户上传的PDF的路径。
        base_work_dir (str): 本次会话的临时工作目录。
        progress (Any | None): Gradio progress object.

    Returns:
        str: 最终配对结果的目录路径，如果失败则返回 None。
    """
    if not all([ImagePDFLoader, extract_and_save_layout_components]):
        tqdm.write("[!] 错误: figure_pipeline 的一个或多个核心依赖项未能加载。")
        return None

    pdf_file = Path(pdf_path)
    pdf_stem = pdf_file.stem
    model_path = "pragent/model/doclayout_yolo_docstructbench_imgsz1024.pt"

    tqdm.write(f"\n--- 步骤 1/3: 将PDF '{pdf_file.name}' 转换为图片 ---")
    page_save_dir = os.path.join(base_work_dir, "page_paper", pdf_stem)
    os.makedirs(page_save_dir, exist_ok=True)
    try:
        loader = ImagePDFLoader(pdf_path)
        page_image_paths = []
        for i, img in enumerate(loader.load()):
            path = os.path.join(page_save_dir, f"page_{i+1}.png")
            img.save(path)
            page_image_paths.append(path)
        tqdm.write(f"[*] 所有 {len(page_image_paths)} 页已保存至: {page_save_dir}")
    except Exception as e:
        tqdm.write(f"[!] 错误：加载或转换PDF时失败: {e}")
        return None

    if len(page_image_paths) > 20:
        tqdm.write(f"[!] Warning: PDF has {len(page_image_paths)} pages. Processing only the first 20 pages to avoid timeout.")
        page_image_paths = page_image_paths[:20]

    tqdm.write(f"\n--- 步骤 2/3: 分析页面布局以裁剪图和表 ---")
    cropped_results_dir = os.path.join(base_work_dir, "cropped_results", pdf_stem)
    num_pages = len(page_image_paths)
    for i, path in enumerate(page_image_paths):
        if progress:
            progress(0.3 + (i / num_pages) * 0.2, desc=f"Analyzing page {i+1}/{num_pages}")
        
        page_num_str = Path(path).stem
        page_crop_dir = os.path.join(cropped_results_dir, page_num_str)
        
        await asyncio.to_thread(
            extract_and_save_layout_components,
            image_path=path,
            model_path=model_path,
            save_base_dir=page_crop_dir,
            imgsz=640
        )

    tqdm.write(f"[*] 所有裁剪结果已保存至: {cropped_results_dir}")

    tqdm.write(f"\n--- 步骤 3/3: 对裁剪出的组件进行配对 ---")
    final_paired_dir = os.path.join(base_work_dir, "paired_results", pdf_stem)
    run_pairing_process(cropped_results_dir, final_paired_dir, threshold=30)
    
    if os.path.isdir(final_paired_dir):
        return final_paired_dir
    return None

def run_pairing_process(source_dir_str: str, output_dir_str: str, threshold: int):
    """配对逻辑，现在是pipeline的一部分。"""
    source_dir = Path(source_dir_str)
    output_root_dir = Path(output_dir_str)
    if output_root_dir.exists(): shutil.rmtree(output_root_dir)
    output_root_dir.mkdir(parents=True, exist_ok=True)
    
    tqdm.write(f"    开始最近邻配对流程 (阈值 = {threshold})")

    page_dirs = sorted([d for d in source_dir.iterdir() if d.is_dir() and d.name.startswith('page_')])
    for page_dir in page_dirs:
        output_page_dir = output_root_dir / page_dir.name
        output_page_dir.mkdir(exist_ok=True)
        pair_items_on_page(str(page_dir), str(output_page_dir), threshold)

def pair_items_on_page(page_dir: str, output_dir: str, threshold: int):
    """处理单个页面目录，进行最近邻配对。"""
    organized_files = defaultdict(dict)
    component_types = ["figure", "figure_caption", "table", "table_caption_above", "table_caption_below"]
    
    def parse_filename(filename: str):
        match = re.match(r'([a-zA-Z_]+)_(\d+)_score([\d.]+)\.jpg', filename)
        return (match.group(1), int(match.group(2))) if match else (None, None)

    for comp_type in component_types:
        comp_dir = os.path.join(page_dir, comp_type)
        if os.path.isdir(comp_dir):
            for filename in os.listdir(comp_dir):
                _, index = parse_filename(filename)
                if index is not None: organized_files[comp_type][index] = os.path.join(comp_dir, filename)

    paired_files, used_captions = set(), defaultdict(set)

    for item_type, cap_types in [("figure", ["figure_caption"]), ("table", ["table_caption_above", "table_caption_below"])]:
        for item_index, item_path in organized_files[item_type].items():
            best_match = {'min_diff': float('inf'), 'cap_path': None, 'cap_index': -1, 'cap_type': ''}
            for cap_type in cap_types:
                for cap_index, cap_path in organized_files[cap_type].items():
                    if cap_index in used_captions[cap_type]: continue
                    diff = abs(item_index - cap_index)
                    if diff < best_match['min_diff']:
                        best_match.update({'min_diff': diff, 'cap_path': cap_path, 'cap_index': cap_index, 'cap_type': cap_type})
            
            if best_match['cap_path'] and best_match['min_diff'] <= threshold:
                target_dir = os.path.join(output_dir, f"paired_{item_type}_{item_index}")
                os.makedirs(target_dir, exist_ok=True)
                shutil.copy(item_path, target_dir); shutil.copy(best_match['cap_path'], target_dir)
                paired_files.add(item_path); paired_files.add(best_match['cap_path'])
                used_captions[best_match['cap_type']].add(best_match['cap_index'])

    for files_dict in organized_files.values():
        for file_path in files_dict.values():
            if file_path not in paired_files:
                item_type, index = parse_filename(Path(file_path).name)
                if item_type:
                    target_dir = os.path.join(output_dir, f"unpaired_{item_type}_{index}")
                    os.makedirs(target_dir, exist_ok=True); shutil.copy(file_path, target_dir)