# batch_processor.py
import multiprocessing
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
import tqdm
from typing import List, Callable
import os
import json


class BatchProcessor:
    """批量处理器"""

    def __init__(self, max_workers=None):
        if max_workers is None:
            max_workers = multiprocessing.cpu_count()
        self.max_workers = max_workers

    def process_pdfs_batch(self, pdf_paths: List[str], process_func: Callable):
        """批量处理PDF"""
        print(f"开始批量处理 {len(pdf_paths)} 个PDF，使用 {self.max_workers} 个工作进程")

        results = []
        with ProcessPoolExecutor(max_workers=self.max_workers) as executor:
            # 提交任务
            future_to_path = {
                executor.submit(process_func, pdf_path): pdf_path
                for pdf_path in pdf_paths
            }

            # 使用进度条
            with tqdm.tqdm(total=len(pdf_paths), desc="处理PDF") as pbar:
                for future in concurrent.futures.as_completed(future_to_path):
                    pdf_path = future_to_path[future]
                    try:
                        result = future.result()
                        if result:
                            results.append(result)
                    except Exception as e:
                        print(f"处理失败 {pdf_path}: {e}")

                    pbar.update(1)

        return results

    def parallel_search(self, queries: List[str], search_func: Callable, **kwargs):
        """并行搜索"""
        with ThreadPoolExecutor(max_workers=min(self.max_workers, 10)) as executor:
            futures = [executor.submit(search_func, query, **kwargs) for query in queries]
            results = [future.result() for future in futures]
        return results


# 使用示例
def batch_process_main():
    processor = BatchProcessor(max_workers=4)

    # 批量处理PDF
    pdf_paths = [f for f in os.listdir('./pdfs') if f.endswith('.pdf')]
    pdf_paths = [os.path.join('./pdfs', f) for f in pdf_paths]

    from pdf_processor import PDFSectionExtractor
    extractor = PDFSectionExtractor()

    results = processor.process_pdfs_batch(
        pdf_paths,
        lambda x: extractor.process_pdf(x)
    )

    # 保存结果
    with open('batch_processed.json', 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=2)