"""
流式处理版本的文档处理流水线

优化点：
1. 按页处理：处理一页，立即保存该页的JSON和Markdown
2. 内存优化：每页处理完成后立即释放该页的图像和OCR数据
3. 输出一致：使用 OutputFormatterV2 的方法，确保输出格式与批量模式完全一致

内存优化：
- 每页处理完成后立即保存并释放内存
- 只保留页面结果用于跨页表格合并
- 最后统一生成完整输出（middle.json、完整Markdown等）
"""
import sys
from typing import Dict, List, Any, Optional
from pathlib import Path
from loguru import logger

# 添加 ocr_platform 根目录到 Python 路径（用于导入 ocr_utils）
ocr_platform_root = Path(__file__).parents[3]  # core -> universal_doc_parser -> ocr_tools -> ocr_platform -> repository.git
if str(ocr_platform_root) not in sys.path:
    sys.path.insert(0, str(ocr_platform_root))

# 添加 universal_doc_parser 根目录到 Python 路径（用于相对导入）
module_root = Path(__file__).parents[0]  # core -> universal_doc_parser
if str(module_root) not in sys.path:
    sys.path.insert(0, str(module_root))

# 导入基础类（复用现有实现）
from .pipeline_manager_v2 import EnhancedDocPipeline
from .pdf_utils import PDFUtils

# 从 ocr_utils 导入输出格式化器
try:
    from ocr_utils import (
        JSONFormatters,
        VisualizationUtils,
        MarkdownGenerator,
        OutputFormatterV2
    )
except ImportError:
    # 降级：尝试从 utils 导入（向后兼容）
    try:
        from ..utils import JSONFormatters, VisualizationUtils, MarkdownGenerator, OutputFormatterV2
    except ImportError:
        from universal_doc_parser.utils import JSONFormatters, VisualizationUtils, MarkdownGenerator, OutputFormatterV2


class StreamingDocPipeline(EnhancedDocPipeline):
    """
    流式处理版本的文档处理流水线
    
    特点：
    - 按页处理，立即保存，释放内存
    - 流式生成Markdown（边处理边写入）
    - 最后统一生成完整Markdown（跨页表格合并）
    """
    
    def __init__(self, config_path: str, output_dir: str):
        """
        初始化流式处理流水线
        
        Args:
            config_path: 配置文件路径
            output_dir: 输出目录（用于立即保存每页结果）
        """
        super().__init__(config_path)
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        # 使用 OutputFormatterV2 进行格式化（保持输出一致）
        self.formatter = OutputFormatterV2(str(self.output_dir))
        
        # 存储已处理的页面（用于跨页表格合并和最终生成）
        self.processed_pages = []
        
        # 存储页面元数据（用于跨页表格合并）
        self.page_metadata = []
        
        # 文档元数据
        self.doc_name = None
        self.is_pdf = False
        self.total_pages = 0
    
    def process_document_streaming(
        self,
        document_path: str,
        page_range: Optional[str] = None,
        output_config: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        流式处理文档主流程
        
        Args:
            document_path: 文档路径（PDF、图片或目录）
            page_range: 页面范围字符串
            output_config: 输出配置
            
        Returns:
            处理结果摘要
        """
        if output_config is None:
            output_config = {}
        
        doc_path = Path(document_path)
        doc_name = doc_path.stem
        
        # 判断输入类型
        is_pdf = doc_path.suffix.lower() == '.pdf'
        
        # 初始化结果摘要
        results_summary = {
            'scene': self.scene_name,
            'document_path': str(doc_path),
            'total_pages': 0,
            'processed_pages': [],
            'metadata': {},
            'output_paths': {
                'json_pages': [],
                'images': [],
            }
        }
        
        try:
            # 1. 加载文档并分类
            dpi = self.config.get('input', {}).get('dpi', 200)
            images, pdf_type, pdf_doc = PDFUtils.load_and_classify_document(
                doc_path, dpi=dpi, page_range=page_range
            )
            
            results_summary['metadata']['pdf_type'] = pdf_type
            results_summary['metadata']['page_count'] = len(images)
            results_summary['metadata']['page_range'] = page_range
            results_summary['total_pages'] = len(images)
            
            # 保存文档元数据
            self.doc_name = doc_name
            self.is_pdf = is_pdf
            self.total_pages = len(images)
            
            logger.info(f"📄 Loaded {len(images)} pages, type: {pdf_type}")
            
            # 创建images子目录
            images_dir = self.output_dir / 'images'
            images_dir.mkdir(exist_ok=True)
            
            # 3. 按页处理（流式）
            for idx, image_dict in enumerate(images):
                page_idx = image_dict.get('page_idx', idx)
                page_name = image_dict.get('page_name', f'page_{page_idx + 1:03d}')
                
                logger.info(f"🔍 Processing page {idx + 1}/{len(images)} (original index: {page_idx})")
                
                # 处理单页
                page_result = self._process_single_page(
                    image_dict=image_dict,
                    page_idx=page_idx,
                    pdf_type=pdf_type,
                    pdf_doc=pdf_doc,
                    page_name=page_name
                )
                
                # 立即保存该页结果（使用 OutputFormatterV2 的方法，保持输出一致）
                page_output_paths = self._save_page_immediately(
                    page_result=page_result,
                    doc_name=doc_name,
                    page_idx=page_idx,
                    is_pdf=is_pdf,
                    images_dir=images_dir,
                    output_config=output_config
                )
                
                # 保存页面到列表（用于跨页表格合并和最终生成，需要在删除前保存）
                self.processed_pages.append(page_result.copy())
                
                # 更新结果摘要
                results_summary['processed_pages'].append({
                    'page_idx': page_idx,
                    'page_name': page_name,
                    'output_paths': page_output_paths
                })
                if page_output_paths.get('json_path'):
                    results_summary['output_paths']['json_pages'].append(page_output_paths['json_path'])
                if page_output_paths.get('image_paths'):
                    results_summary['output_paths']['images'].extend(page_output_paths['image_paths'])
                
                # 保存页面元数据（用于跨页表格合并）
                self.page_metadata.append({
                    'page_idx': page_idx,
                    'page_name': page_name,
                    'json_path': page_output_paths.get('json_path'),
                    'has_table': any(
                        elem.get('type') in ['table', 'table_body']
                        for elem in page_result.get('elements', [])
                    )
                })
                
                # 释放该页的内存（显式删除大对象，但保留副本在 processed_pages 中）
                del page_result
                del image_dict
                
                logger.info(f"✅ Page {idx + 1} processed and saved")
            
            # 4. 构建完整结果（用于跨页表格合并和最终生成）
            complete_results = {
                'scene': self.scene_name,
                'document_path': str(doc_path),
                'pages': self.processed_pages,
                'metadata': results_summary['metadata']
            }
            
            # 5. 跨页表格合并
            if output_config.get('merge_cross_page_tables', True):
                logger.info("🔄 Merging cross-page tables...")
                complete_results = PDFUtils.merge_cross_page_tables(complete_results)
            
            # 6. 使用 OutputFormatterV2 生成所有输出（保持输出一致）
            logger.info("💾 Saving all results using OutputFormatterV2...")
            final_output_paths = self.formatter.save_results(complete_results, output_config)
            
            # 更新结果摘要
            results_summary['output_paths'].update(final_output_paths)
            
            # 7. 关闭 PDF 文档
            if pdf_doc is not None:
                pdf_doc.close()
            
            logger.info(f"✅ Document processing completed (streaming mode)")
            return results_summary
            
        except Exception as e:
            logger.error(f"❌ Document processing failed: {e}")
            raise
    
    def _save_page_immediately(
        self,
        page_result: Dict[str, Any],
        doc_name: str,
        page_idx: int,
        is_pdf: bool,
        images_dir: Path,
        output_config: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        立即保存单页结果（使用 OutputFormatterV2 的方法，保持输出一致）
        
        Returns:
            保存的文件路径字典
        """
        output_paths = {}
        
        # 构造单页 results 格式（用于调用 OutputFormatterV2 的方法）
        single_page_results = {
            'document_path': '',  # 不需要完整路径
            'pages': [page_result],
            'metadata': {}
        }
        
        # 1. 保存图片元素（使用 VisualizationUtils）
        if output_config.get('save_images', True):
            image_paths = VisualizationUtils.save_image_elements(
                single_page_results,
                images_dir,
                doc_name,
                is_pdf=is_pdf
            )
            if image_paths:
                output_paths['image_paths'] = image_paths
        
        # 2. 保存单页JSON（使用 JSONFormatters.save_page_jsons，保持命名规则一致）
        if output_config.get('save_page_json', True):
            normalize_numbers = output_config.get('normalize_numbers', True)
            page_json_paths = JSONFormatters.save_page_jsons(
                single_page_results,
                self.output_dir,
                doc_name,
                is_pdf=is_pdf,
                normalize_numbers=normalize_numbers
            )
            if page_json_paths:
                output_paths['json_path'] = page_json_paths[0]  # 单页只有一个JSON
        
        # 3. 保存单页Markdown（使用 MarkdownGenerator.save_page_markdowns，保持命名规则一致）
        if output_config.get('save_page_markdown', True):
            normalize_numbers = output_config.get('normalize_numbers', True)
            page_md_paths = MarkdownGenerator.save_page_markdowns(
                single_page_results,
                self.output_dir,
                doc_name,
                is_pdf=is_pdf,
                normalize_numbers=normalize_numbers
            )
            if page_md_paths:
                output_paths['markdown_path'] = page_md_paths[0]  # 单页只有一个Markdown
        
        # 4. 保存可视化图片（如果开启debug）
        if output_config.get('save_layout_image', False):
            VisualizationUtils.save_layout_images(
                single_page_results,
                self.output_dir,
                doc_name,
                is_pdf=is_pdf
            )
        
        if output_config.get('save_ocr_image', False):
            VisualizationUtils.save_ocr_images(
                single_page_results,
                self.output_dir,
                doc_name,
                is_pdf=is_pdf
            )
        
        return output_paths
    

