"""
增强版文档处理流水线 v2

按照优化后的流程实现：
1. PDF分类 → 扫描件/数字原生PDF
2. 页面方向识别（仅扫描件）
3. Layout检测
4. 并行处理：
   - 文本区域：OCR检测+识别 / PDF字符提取
   - 表格区域：OCR检测（坐标） + VLM结构识别 → 坐标匹配
5. 合并结果 → 跨页表格合并

模块结构：
- coordinate_utils.py: 坐标转换工具
- layout_utils.py: 布局处理工具
- pdf_utils.py: PDF处理工具
- element_processors.py: 元素处理器
"""
import sys
from typing import Dict, List, Any, Optional
from pathlib import Path
import numpy as np
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))

# 导入本地模块
try:
    from .model_factory import ModelFactory
    from .config_manager import ConfigManager
    from .coordinate_utils import CoordinateUtils
    from .layout_utils import LayoutUtils, SpanMatcher
    from .pdf_utils import PDFUtils
    from .element_processors import ElementProcessors
except ImportError:
    from model_factory import ModelFactory
    from config_manager import ConfigManager
    from coordinate_utils import CoordinateUtils
    from layout_utils import LayoutUtils, SpanMatcher
    from pdf_utils import PDFUtils
    from element_processors import ElementProcessors

# 从 ocr_tools.ocr_merger 导入 merger 组件
try:
    from ocr_tools.ocr_merger import TableCellMatcher, TextMatcher
    MERGER_AVAILABLE = True
except ImportError:
    MERGER_AVAILABLE = False
    TableCellMatcher = None
    TextMatcher = None


class EnhancedDocPipeline:
    """增强版文档处理流水线"""
    
    # ==================== 元素类别定义 ====================
    # 参考 MinerU CategoryId 定义
    
    # 文本类元素：需要OCR识别或PDF文本提取
    TEXT_CATEGORIES = [
        'text', 'title', 'ocr_text', 'low_score_text',  # 基础文本
        'header', 'footer', 'page_number',               # 页面元素
        'ref_text', 'aside_text', 'page_footnote',       # 辅助文本
    ]
    
    # 表格相关元素
    TABLE_BODY_CATEGORIES = ['table', 'table_body']
    TABLE_TEXT_CATEGORIES = ['table_caption', 'table_footnote']
    
    # 图片相关元素
    IMAGE_BODY_CATEGORIES = ['image', 'image_body', 'figure']
    IMAGE_TEXT_CATEGORIES = ['image_caption', 'image_footnote']
    
    # 公式类元素
    EQUATION_CATEGORIES = [
        'interline_equation', 'inline_equation', 'equation',
        'interline_equation_yolo', 'interline_equation_number'
    ]
    
    # 丢弃类元素（水印、装饰等）
    DISCARD_CATEGORIES = ['abandon', 'discarded']
    
    # 代码类元素
    CODE_CATEGORIES = ['code', 'code_body', 'code_caption', 'algorithm']
    
    # ==================== 初始化 ====================
    
    def __init__(self, config_path: str):
        """
        初始化流水线
        
        Args:
            config_path: 配置文件路径
        """
        self.config = ConfigManager.load_config(config_path)
        self.scene_name = self.config.get('scene_name', 'unknown')
        self.debug_mode = self.config.get('output', {}).get('debug_mode', False)
        
        # 初始化组件
        self._init_components()
        
        # 初始化元素处理器
        self._init_element_processors()
        
        logger.info(f"✅ Pipeline initialized for scene: {self.scene_name}")
    
    def _init_components(self):
        """初始化处理组件"""
        try:
            # 1. 预处理器（方向分类）
            self.preprocessor = ModelFactory.create_preprocessor(
                self.config['preprocessor']
            )
            
            # 2. 版式检测器
            self.layout_detector = ModelFactory.create_layout_detector(
                self.config['layout_detection']
            )
            
            # 3. VL识别器（表格、公式）
            if self.config.get('vl_recognition', {}) != {}:
                self.vl_recognizer = ModelFactory.create_vl_recognizer(
                    self.config['vl_recognition']
                )
            else:
                self.vl_recognizer = None
            
            # 4. OCR识别器
            self.ocr_recognizer = ModelFactory.create_ocr_recognizer(
                self.config['ocr_recognition']
            )

            # 5. 有线表格识别器（可选）
            self.table_config = self.config.get('table_recognition_wired', {})
            self.wired_table_recognizer = None
            if self.table_config.get('use_wired_unet', False):
                try:
                    from models.adapters.mineru_wired_table import MinerUWiredTableRecognizer
                    ocr_engine = getattr(self.ocr_recognizer, 'ocr_model', None)
                    self.wired_table_recognizer = MinerUWiredTableRecognizer(
                        self.table_config, ocr_engine
                    )
                    logger.info("✅ Wired UNet table recognizer initialized")
                except Exception as e:
                    logger.warning(f"⚠️ Wired UNet table recognizer init failed: {e}")
            
        except Exception as e:
            logger.error(f"❌ Failed to initialize pipeline: {e}")
            raise
    
    def _init_element_processors(self):
        """初始化元素处理器"""
        # 初始化 merger 组件
        table_cell_matcher = None
        if MERGER_AVAILABLE and TextMatcher is not None and TableCellMatcher is not None:
            text_matcher = TextMatcher()
            table_cell_matcher = TableCellMatcher(
                text_matcher=text_matcher,
                x_tolerance=3,
                y_tolerance=10
            )
        else:
            logger.warning("⚠️ Merger components not available, cell coordinate matching disabled")
        
        # 创建元素处理器
        self.element_processors = ElementProcessors(
            preprocessor=self.preprocessor,
            ocr_recognizer=self.ocr_recognizer,
            vl_recognizer=self.vl_recognizer,
            table_cell_matcher=table_cell_matcher,
            wired_table_recognizer=getattr(self, 'wired_table_recognizer', None),
        )
    
    # ==================== 主处理流程 ====================
    
    def process_document(
        self, 
        document_path: str,
        page_range: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        处理文档主流程
        
        Args:
            document_path: 文档路径（PDF、图片或目录）
            page_range: 页面范围字符串，如 "1-5,7,9-12"
                       - PDF：按页码（从1开始）
                       - 图片目录：按文件名排序后的位置（从1开始）
            
        Returns:
            处理结果字典
        """
        doc_path = Path(document_path)
        
        results = {
            'scene': self.scene_name,
            'document_path': str(doc_path),
            'pages': [],
            'metadata': {},
            'debug_info': {} if self.debug_mode else None
        }
        
        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['metadata']['pdf_type'] = pdf_type
            results['metadata']['page_count'] = len(images)
            results['metadata']['page_range'] = page_range
            
            logger.info(f"📄 Loaded {len(images)} pages, type: {pdf_type}")
            
            # 2. 处理每一页
            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
                )
                results['pages'].append(page_result)
            
            # 3. 跨页表格合并
            results = PDFUtils.merge_cross_page_tables(results)
            
            # 4. 关闭 PDF 文档
            if pdf_doc is not None:
                pdf_doc.close()
            
            logger.info(f"✅ Document processing completed")
            return results
            
        except Exception as e:
            logger.error(f"❌ Document processing failed: {e}")
            raise
    
    def _process_single_page(
        self,
        image_dict: Dict[str, Any],
        page_idx: int,
        pdf_type: str,
        pdf_doc: Optional[Any] = None,
        page_name: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        处理单页文档
        
        流程：
        1. 扫描件：页面方向识别 → Layout检测
        2. 数字PDF：直接Layout检测
        3. 并行处理文本区域和表格区域
        4. 合并结果
        5. 坐标转换回原始图片坐标系
        """
        pil_image = image_dict['img_pil']
        scale = image_dict.get('scale', 1.0)
        original_image = np.array(pil_image)
        
        # 页面名称（用于输出文件命名）
        if page_name is None:
            page_name = f"page_{page_idx + 1:03d}"
        
        page_result = {
            'page_idx': page_idx,
            'page_name': page_name,  # 用于输出文件命名
            'elements': [],
            'image_shape': original_image.shape,
            'original_image': original_image,
            'angle': 0,
            'scale': scale,
            'pdf_type': pdf_type
        }
        
        # 用于检测的图片(可能被旋转)
        detection_image = original_image.copy()
        rotate_angle = 0
        
        # 1. 页面方向识别（仅扫描件）
        if pdf_type == 'ocr':
            try:
                detection_image, rotate_angle = self.preprocessor.process(original_image)
                page_result['angle'] = rotate_angle
                
                if rotate_angle != 0:
                    logger.info(f"📐 Page {page_idx}: rotated {rotate_angle}° for detection")
            except Exception as e:
                logger.warning(f"⚠️ Orientation detection failed: {e}")
        
        # 2. Layout检测
        try:
            layout_results = self.layout_detector.detect(detection_image)
            logger.info(f"📋 Page {page_idx}: detected {len(layout_results)} elements (before dedup)")
        except Exception as e:
            logger.error(f"❌ Layout detection failed: {e}")
            layout_results = []
        
        # 2.5 处理重叠框
        if layout_results:
            original_count = len(layout_results)
            layout_results = LayoutUtils.remove_overlapping_boxes(layout_results)
            removed_count = original_count - len(layout_results)
            if removed_count > 0:
                logger.info(f"🔄 Page {page_idx}: removed {removed_count} overlapping boxes")
        
        # 2.6 将大面积文本块转换为表格（后处理）
        if layout_results:
            convert_large_text = self.config.get('layout', {}).get('convert_large_text_to_table', False)
            if convert_large_text:
                # 获取图像尺寸 (height, width)
                img_shape = detection_image.shape
                if len(img_shape) >= 2:
                    image_shape = (int(img_shape[0]), int(img_shape[1]))
                else:
                    image_shape = (0, 0)
                
                layout_results = LayoutUtils.convert_large_text_to_table(
                    layout_results,
                    image_shape,
                    min_area_ratio=self.config.get('layout', {}).get('min_text_area_ratio', 0.25),
                    min_width_ratio=self.config.get('layout', {}).get('min_text_width_ratio', 0.4),
                    min_height_ratio=self.config.get('layout', {}).get('min_text_height_ratio', 0.3)
                )
        
        page_result['layout_raw'] = layout_results
        
        # 3. 整页 OCR 获取所有 text spans（关键改进）
        all_ocr_spans = []
        try:
            all_ocr_spans = self.ocr_recognizer.recognize_text(detection_image)
            # 去除重复 spans
            all_ocr_spans = SpanMatcher.remove_duplicate_spans(all_ocr_spans)
            # 按坐标排序（从上到下，从左到右），方便人工检查缺失字符
            all_ocr_spans = self._sort_spans_by_position(all_ocr_spans)
            logger.info(f"📝 Page {page_idx}: OCR detected {len(all_ocr_spans)} text spans")
        except Exception as e:
            logger.warning(f"⚠️ Full-page OCR failed: {e}")
        
        # 4. 将 OCR spans 匹配到 layout blocks
        matched_spans = SpanMatcher.match_spans_to_blocks(
            all_ocr_spans, layout_results, overlap_threshold=0.5
        )
        
        # 5. 分类元素
        classified_elements = self._classify_elements(layout_results, page_idx)
        
        # 6. 处理各类元素（传入匹配的 spans）
        processed_elements, discarded_elements = self._process_all_elements(
            detection_image=detection_image,
            classified_elements=classified_elements,
            pdf_type=pdf_type,
            pdf_doc=pdf_doc,
            page_idx=page_idx,
            scale=scale,
            matched_spans=matched_spans,
            layout_results=layout_results
        )
        
        # 7. 按阅读顺序排序
        sorted_elements = LayoutUtils.sort_elements_by_reading_order(processed_elements)
        sorted_discarded = LayoutUtils.sort_elements_by_reading_order(discarded_elements)
        
        # 8. 坐标转换回原始图片坐标系
        if rotate_angle != 0:
            sorted_elements = [
                CoordinateUtils.transform_coords_to_original(
                    element, rotate_angle, detection_image.shape, original_image.shape
                ) for element in sorted_elements
            ]
            sorted_discarded = [
                CoordinateUtils.transform_coords_to_original(
                    element, rotate_angle, detection_image.shape, original_image.shape
                ) for element in sorted_discarded
            ]
            logger.debug(f"📐 Coordinates transformed back to original image")
        
        page_result['elements'] = sorted_elements
        page_result['discarded_blocks'] = sorted_discarded
        return page_result
    
    # ==================== OCR Spans 排序 ====================
    
    @staticmethod
    def _sort_spans_by_position(spans: List[Dict[str, Any]], row_tolerance: float = 10.0) -> List[Dict[str, Any]]:
        """
        按照坐标对OCR spans排序（从上到下，从左到右）
        
        Args:
            spans: OCR检测的text spans列表
            row_tolerance: 同一行的y坐标容差（像素）
            
        Returns:
            排序后的spans列表
        """
        if not spans:
            return spans
        
        def get_sort_key(span: Dict[str, Any]) -> tuple:
            """
            生成排序键
            
            Returns:
                (row_group, x1) - row_group用于分组同一行，x1用于行内从左到右排序
            """
            bbox = span.get('bbox', [[0, 0], [0, 0], [0, 0], [0, 0]])
            
            # 获取左上角坐标
            if isinstance(bbox[0], (list, tuple)):
                # 四点格式 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
                x1, y1 = bbox[0][0], bbox[0][1]
            else:
                # [x1, y1, x2, y2] 格式
                x1, y1 = bbox[0], bbox[1]
            
            # 将y坐标按row_tolerance分组（同一行）
            row_group = int(y1 / row_tolerance)
            
            return (row_group, x1)
        
        # 执行排序
        sorted_spans = sorted(spans, key=get_sort_key)
        return sorted_spans
    
    # ==================== 元素分类与处理 ====================
    
    def _classify_elements(
        self,
        layout_results: List[Dict[str, Any]],
        page_idx: int
    ) -> Dict[str, List[Dict[str, Any]]]:
        """
        将布局结果分类为不同类型的元素
        
        Args:
            layout_results: 布局检测结果
            page_idx: 页码
            
        Returns:
            分类后的元素字典
        """
        classified = {
            'text': [],
            'table_body': [],
            'table_text': [],
            'image_body': [],
            'image_text': [],
            'equation': [],
            'code': [],
            'discard': []
        }
        
        for item in layout_results:
            category = item.get('category', '').lower()
            
            if category in self.TEXT_CATEGORIES:
                classified['text'].append(item)
            elif category in self.TABLE_BODY_CATEGORIES:
                classified['table_body'].append(item)
            elif category in self.TABLE_TEXT_CATEGORIES:
                classified['table_text'].append(item)
            elif category in self.IMAGE_BODY_CATEGORIES:
                classified['image_body'].append(item)
            elif category in self.IMAGE_TEXT_CATEGORIES:
                classified['image_text'].append(item)
            elif category in self.EQUATION_CATEGORIES:
                classified['equation'].append(item)
            elif category in self.CODE_CATEGORIES:
                classified['code'].append(item)
            elif category in self.DISCARD_CATEGORIES:
                classified['discard'].append(item)
            else:
                logger.debug(f"Unknown category '{category}', treating as text")
                classified['text'].append(item)
        
        logger.info(f"📊 Page {page_idx} elements: "
                   f"text={len(classified['text'])}, "
                   f"table={len(classified['table_body'])}, "
                   f"table_text={len(classified['table_text'])}, "
                   f"image={len(classified['image_body'])}, "
                   f"image_text={len(classified['image_text'])}, "
                   f"equation={len(classified['equation'])}, "
                   f"code={len(classified['code'])}, "
                   f"discard={len(classified['discard'])}")
        
        return classified
    
    def _process_all_elements(
        self,
        detection_image: np.ndarray,
        classified_elements: Dict[str, List[Dict[str, Any]]],
        pdf_type: str,
        pdf_doc: Optional[Any],
        page_idx: int,
        scale: float,
        matched_spans: Optional[Dict[int, List[Dict[str, Any]]]] = None,
        layout_results: Optional[List[Dict[str, Any]]] = None
    ) -> tuple:
        """
        处理所有分类后的元素
        
        Args:
            detection_image: 检测用图像
            classified_elements: 分类后的元素
            pdf_type: PDF类型
            pdf_doc: PDF文档对象
            page_idx: 页码
            scale: 缩放比例
            matched_spans: 匹配的 OCR spans {block_idx: [spans]}
            layout_results: 原始 layout 检测结果（用于索引匹配）
            
        Returns:
            (processed_elements, discarded_elements)
        """
        processed_elements = []
        discarded_elements = []
        
        # 构建 layout item 到 block index 的映射
        item_to_block_idx = {}
        if layout_results:
            for idx, item in enumerate(layout_results):
                # 使用 bbox 作为唯一标识（转为元组方便哈希）
                bbox_key = tuple(item.get('bbox', [0, 0, 0, 0]))
                item_to_block_idx[bbox_key] = idx
        
        def get_matched_spans_for_item(item: Dict[str, Any]) -> List[Dict[str, Any]]:
            """获取某个 layout item 匹配的 spans"""
            if not matched_spans or not layout_results:
                return []
            bbox_key = tuple[Any, ...](item.get('bbox', [0, 0, 0, 0]))
            block_idx = item_to_block_idx.get(bbox_key, -1)
            return matched_spans.get(block_idx, [])
        
        # 处理纯文本区域
        for item in classified_elements['text']:
            try:
                spans = get_matched_spans_for_item(item)
                element = self.element_processors.process_text_element(
                    detection_image, item, pdf_type, pdf_doc, page_idx, scale,
                    pre_matched_spans=spans
                )
                processed_elements.append(element)
            except Exception as e:
                logger.warning(f"⚠️ Text processing failed: {e}")
                processed_elements.append(ElementProcessors.create_error_element(item, str(e)))
        
        # 处理表格标题/脚注
        for item in classified_elements['table_text']:
            try:
                spans = get_matched_spans_for_item(item)
                element = self.element_processors.process_text_element(
                    detection_image, item, pdf_type, pdf_doc, page_idx, scale,
                    pre_matched_spans=spans
                )
                processed_elements.append(element)
            except Exception as e:
                logger.warning(f"⚠️ Table text processing failed: {e}")
                processed_elements.append(ElementProcessors.create_error_element(item, str(e)))
        
        # 处理图片标题/脚注
        for item in classified_elements['image_text']:
            try:
                spans = get_matched_spans_for_item(item)
                element = self.element_processors.process_text_element(
                    detection_image, item, pdf_type, pdf_doc, page_idx, scale,
                    pre_matched_spans=spans
                )
                processed_elements.append(element)
            except Exception as e:
                logger.warning(f"⚠️ Image text processing failed: {e}")
                processed_elements.append(ElementProcessors.create_error_element(item, str(e)))
        
        # 处理表格主体
        for item in classified_elements['table_body']:
            try:
                spans = get_matched_spans_for_item(item)
                
                # 🔑 关键：根据配置选择表格识别路径
                use_wired_unet = self.table_config.get('use_wired_unet', False)
                
                if use_wired_unet and self.wired_table_recognizer:
                    # 有线表格路径：UNet 识别
                    logger.info(f"🔷 Using wired UNet table recognition (configured)")
                    element = self.element_processors.process_table_element_wired(
                        detection_image, item, scale, pre_matched_spans=spans
                    )
                    
                    # 如果有线识别失败（返回空 HTML），fallback 到 VLM
                    if not element['content'].get('html') and not element['content'].get('cells'):
                        raise ValueError(f"Wired UNet table recognition failed, element: {item}")
                else:
                    # VLM 无线表格路径（默认）
                    logger.info(f"🔷 Using VLM table recognition (configured)")
                    element = self.element_processors.process_table_element_vlm(
                        detection_image, item, scale, pre_matched_spans=spans
                    )
                
                processed_elements.append(element)
            except Exception as e:
                logger.warning(f"⚠️ Table processing failed: {e}")
                processed_elements.append(ElementProcessors.create_error_element(item, str(e)))
        
        # 处理公式元素
        for item in classified_elements['equation']:
            try:
                element = self.element_processors.process_equation_element(
                    detection_image, item
                )
                processed_elements.append(element)
            except Exception as e:
                logger.warning(f"⚠️ Equation processing failed: {e}")
                processed_elements.append(ElementProcessors.create_error_element(item, str(e)))
        
        # 处理图片主体
        for item in classified_elements['image_body']:
            try:
                element = self.element_processors.process_image_element(
                    detection_image, item
                )
                processed_elements.append(element)
            except Exception as e:
                logger.warning(f"⚠️ Image processing failed: {e}")
                processed_elements.append(ElementProcessors.create_error_element(item, str(e)))
        
        # 处理代码元素
        for item in classified_elements['code']:
            try:
                spans = get_matched_spans_for_item(item)
                element = self.element_processors.process_code_element(
                    detection_image, item, pdf_type, pdf_doc, page_idx, scale,
                    pre_matched_spans=spans
                )
                processed_elements.append(element)
            except Exception as e:
                logger.warning(f"⚠️ Code processing failed: {e}")
                processed_elements.append(ElementProcessors.create_error_element(item, str(e)))
        
        # 处理丢弃元素
        for item in classified_elements['discard']:
            try:
                spans = get_matched_spans_for_item(item)
                element = self.element_processors.process_discard_element(
                    detection_image, item, pre_matched_spans=spans
                )
                discarded_elements.append(element)
            except Exception as e:
                logger.debug(f"Discard element processing failed: {e}")
        
        return processed_elements, discarded_elements
    
    # ==================== 资源管理 ====================
    
    def cleanup(self):
        """清理资源"""
        try:
            if hasattr(self, 'preprocessor'):
                self.preprocessor.cleanup()
            if hasattr(self, 'layout_detector'):
                self.layout_detector.cleanup()
            if hasattr(self, 'vl_recognizer'):
                self.vl_recognizer.cleanup()
            if hasattr(self, 'ocr_recognizer'):
                self.ocr_recognizer.cleanup()
            logger.info("✅ Pipeline cleanup completed")
        except Exception as e:
            logger.warning(f"⚠️ Cleanup failed: {e}")
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.cleanup()
