#!/usr/bin/env python3
"""PaddleOCR VL Python API 包装脚本
用于独立进程调用 PaddleOCR，避免与主进程共享显存
"""

import os
import sys
import argparse
from pathlib import Path


def main():
    parser = argparse.ArgumentParser(description='PaddleOCR VL Python API 包装脚本')
    parser.add_argument('-i', '--input', required=True, help='输入图片路径')
    parser.add_argument('-o', '--output', required=True, help='输出目录路径')
    parser.add_argument('--mode', choices=['ocr', 'doc_parser'], default='doc_parser',
                        help='处理模式: ocr(纯文本) 或 doc_parser(图表识别)')
    parser.add_argument('--device', default='gpu:0', help='设备类型，如 gpu:0 或 cpu')
    parser.add_argument('--use-chart', action='store_true', help='启用图表识别')
    parser.add_argument('--use-layout', action='store_true', help='启用版面检测')
    
    args = parser.parse_args()
    
    # 检查输入文件
    if not os.path.exists(args.input):
        print(f"错误: 输入文件不存在: {args.input}", file=sys.stderr)
        sys.exit(1)
    
    # 创建输出目录
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    try:
        from paddleocr import PaddleOCRVL
        
        # 根据模式设置参数
        if args.mode == 'ocr':
            # 纯文本识别模式
            use_chart = False
            use_layout = False
        else:
            # 图表识别模式
            use_chart = args.use_chart
            use_layout = args.use_layout
        
        print(f"初始化 PaddleOCR VL (mode={args.mode}, chart={use_chart}, layout={use_layout})...", file=sys.stderr)
        
        # 初始化 pipeline
        pipeline = PaddleOCRVL(
            device=args.device,
            use_doc_unwarping=False,
            use_doc_orientation_classify=True,
            use_chart_recognition=use_chart,
            use_layout_detection=use_layout,
        )
        
        print(f"开始处理: {args.input}", file=sys.stderr)
        
        # 执行识别
        result = pipeline.predict(input=args.input)
        
        # 保存结果
        image_basename = os.path.splitext(os.path.basename(args.input))[0]
        
        for item in result:
            if hasattr(item, 'save'):
                item_save_path = output_dir / image_basename
                item.save(str(item_save_path))
                print(f"结果已保存到: {item_save_path}", file=sys.stderr)
                
                # 查找生成的 markdown 文件（多种可能的路径）
                possible_paths = [
                    item_save_path / f"{image_basename}.md",
                    output_dir / f"{image_basename}.md",
                    item_save_path / "result.md",
                ]
                
                for markdown_file in possible_paths:
                    if markdown_file.exists():
                        print(f"SUCCESS: {markdown_file}")
                        sys.exit(0)
                
                # 如果没找到，尝试递归查找所有 .md 文件
                import glob
                md_files = list(item_save_path.rglob("*.md"))
                if md_files:
                    print(f"SUCCESS: {md_files[0]}")
                    sys.exit(0)
        
        print("错误: 未找到 markdown 输出文件", file=sys.stderr)
        print(f"输出目录: {output_dir}", file=sys.stderr)
        print(f"保存路径: {item_save_path if 'item_save_path' in locals() else 'N/A'}", file=sys.stderr)
        sys.exit(1)
        
    except ImportError as e:
        print(f"错误: 无法导入 PaddleOCR 模块: {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"错误: {type(e).__name__}: {e}", file=sys.stderr)
        import traceback
        traceback.print_exc(file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()