"""
元素处理器模块

提供各类文档元素的处理功能：
- 文本元素处理
- 表格元素处理
- 图片元素处理
- 公式元素处理
- 代码元素处理
- 丢弃元素处理
"""
from typing import Dict, List, Any, Optional, Tuple
import numpy as np
from loguru import logger

from .coordinate_utils import CoordinateUtils
from .pdf_utils import PDFUtils

# 导入 SpanMatcher（用于 spans 合并）
try:
    from .layout_utils import SpanMatcher
except ImportError:
    from layout_utils import SpanMatcher

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


class ElementProcessors:
    """元素处理器类"""
    
    def __init__(
        self,
        preprocessor: Any,
        ocr_recognizer: Any,
        vl_recognizer: Any,
        table_cell_matcher: Optional[Any] = None,
        wired_table_recognizer: Optional[Any] = None,
    ):
        """
        初始化元素处理器
        
        Args:
            preprocessor: 预处理器（方向检测）
            ocr_recognizer: OCR识别器
            vl_recognizer: VL识别器（表格、公式）
            table_cell_matcher: 表格单元格匹配器
            wired_table_recognizer: 有线表格识别器（可选）
        """
        self.preprocessor = preprocessor
        self.ocr_recognizer = ocr_recognizer
        self.vl_recognizer = vl_recognizer
        self.table_cell_matcher = table_cell_matcher
        self.wired_table_recognizer = wired_table_recognizer
    
    def _convert_ocr_details_to_absolute(
        self,
        ocr_details: List[Dict[str, Any]],
        region_bbox: List[float]
    ) -> List[Dict[str, Any]]:
        """
        将 OCR 详情中的相对坐标转换为绝对坐标
        
        Args:
            ocr_details: OCR 结果列表，每项包含 'bbox' 字段
            region_bbox: 裁剪区域的绝对坐标 [x1, y1, x2, y2]
            
        Returns:
            坐标转换后的 OCR 详情列表
        """
        if not ocr_details or not region_bbox or len(region_bbox) < 2:
            return ocr_details
        
        converted_details = []
        for item in ocr_details:
            new_item = item.copy()
            ocr_bbox = item.get('bbox', [])
            if ocr_bbox:
                new_item['bbox'] = CoordinateUtils.convert_to_absolute_coords(
                    ocr_bbox, region_bbox
                )
            converted_details.append(new_item)
        
        return converted_details
    
    def process_text_element(
        self,
        image: np.ndarray,
        layout_item: Dict[str, Any],
        pdf_type: str,
        pdf_doc: Optional[Any],
        page_idx: int,
        scale: float,
        pre_matched_spans: Optional[List[Dict[str, Any]]] = None
    ) -> Dict[str, Any]:
        """
        处理文本元素
        
        处理优先级：
        1. 如果有预匹配的 spans（整页 OCR 结果），优先使用
        2. 数字PDF：尝试 PDF 字符提取
        3. 扫描件或提取失败：裁剪区域 OCR（兜底方案）
        
        Args:
            image: 页面图像
            layout_item: 布局检测项
            pdf_type: PDF类型 ('ocr' 或 'txt')
            pdf_doc: PDF文档对象
            page_idx: 页码索引
            scale: 缩放比例
            pre_matched_spans: 预匹配的 OCR spans（来自整页 OCR）
            
        Returns:
            处理后的元素字典
        """
        bbox = layout_item.get('bbox', [0, 0, 0, 0])
        
        text_content = ""
        ocr_details = []
        extraction_method = "none"
        
        # 优先级1：使用预匹配的 spans（整页 OCR 结果）
        if pre_matched_spans and len(pre_matched_spans) > 0:
            text_content, sorted_spans = SpanMatcher.merge_spans_to_text(
                pre_matched_spans, bbox
            )
            if text_content.strip():
                # spans 的坐标已经是绝对坐标，直接使用
                ocr_details = sorted_spans
                extraction_method = "fullpage_ocr"
                logger.debug(f"📝 Text from full-page OCR: '{text_content[:30]}...'")
        
        # 优先级2：数字 PDF 字符提取
        if not text_content.strip() and pdf_type == 'txt' and pdf_doc is not None:
            try:
                text_content, extraction_success = PDFUtils.extract_text_from_pdf(
                    pdf_doc, page_idx, bbox, scale
                )
                if extraction_success and text_content.strip():
                    extraction_method = "pdf_extract"
                    logger.debug(f"📝 Text extracted from PDF: '{text_content[:30]}...'")
            except Exception as e:
                logger.debug(f"PDF text extraction failed: {e}")
        
        # 优先级3：裁剪区域 OCR（兜底方案）
        if not text_content.strip():
            try:
                cropped_image = CoordinateUtils.crop_region(image, bbox)
                ocr_results = self.ocr_recognizer.recognize_text(cropped_image)
                if ocr_results:
                    text_parts = [
                        item['text'] for item in ocr_results 
                        if item.get('confidence', 0) > 0.5
                    ]
                    text_content = " ".join(text_parts)
                    # 将 OCR 坐标转换为绝对坐标
                    ocr_details = self._convert_ocr_details_to_absolute(ocr_results, bbox)
                    extraction_method = "cropped_ocr"
            except Exception as e:
                logger.warning(f"OCR recognition failed: {e}")
        
        return {
            'type': layout_item.get('category', 'text'),
            'bbox': bbox,
            'confidence': layout_item.get('confidence', 0.0),
            'content': {
                'text': text_content,
                'ocr_details': ocr_details,
                'extraction_method': extraction_method
            }
        }
    
    def _prepare_table_ocr(
        self,
        image: np.ndarray,
        bbox: List[float],
        pre_matched_spans: Optional[List[Dict[str, Any]]] = None
    ) -> Tuple[np.ndarray, List[Dict[str, Any]], int, str]:
        """
        表格OCR预处理（共享逻辑）
        
        优先使用整页 OCR 结果，兜底裁剪 OCR
        
        Args:
            image: 页面图像（原始未旋转）
            bbox: 表格的绝对坐标 bbox
            pre_matched_spans: 预匹配的 OCR spans
            
        Returns:
            (cropped_table, ocr_boxes, table_angle, ocr_source)
            其中 cropped_table 已经过方向检测和旋转处理
        """
        cropped_table = CoordinateUtils.crop_region(image, bbox)
        table_angle = 0
        
        # 1. 表格方向检测
        try:
            rotated_table, table_angle = self.preprocessor.process(cropped_table)
            if table_angle != 0:
                logger.info(f"📐 Table rotated {table_angle}°")
                cropped_table = rotated_table
        except Exception as e:
            logger.debug(f"Table orientation detection skipped: {e}")
        
        # 2. 收集 OCR 框：优先整页 OCR，兜底裁剪 OCR
        ocr_boxes = []
        ocr_source = "none"
        
        if pre_matched_spans and len(pre_matched_spans) > 0 and table_angle == 0:
            # 使用整页 OCR 的结果
            for idx, span in enumerate(pre_matched_spans):
                span_bbox = span.get('bbox', [])
                if span_bbox:
                    relative_bbox = [
                        span_bbox[0] - bbox[0],
                        span_bbox[1] - bbox[1],
                        span_bbox[2] - bbox[0],
                        span_bbox[3] - bbox[1]
                    ]
                    formatted_box = CoordinateUtils.convert_ocr_to_matcher_format(
                        relative_bbox,
                        span.get('text', ''),
                        span.get('confidence', 0.0),
                        idx,
                        table_bbox=None
                    )
                    if formatted_box:
                        ocr_boxes.append(formatted_box)
            
            if ocr_boxes:
                ocr_source = "fullpage_ocr"
                logger.info(f"📊 Using {len(ocr_boxes)} text boxes from full-page OCR")
        
        # 兜底：裁剪后单独 OCR
        if not ocr_boxes:
            try:
                ocr_results = self.ocr_recognizer.recognize_text(cropped_table)
                if ocr_results:
                    for idx, item in enumerate(ocr_results):
                        ocr_poly = item.get('bbox', [])
                        if ocr_poly:
                            formatted_box = CoordinateUtils.convert_ocr_to_matcher_format(
                                ocr_poly, 
                                item.get('text', ''),
                                item.get('confidence', 0.0),
                                idx,
                                table_bbox=None
                            )
                            if formatted_box:
                                ocr_boxes.append(formatted_box)
                                # 归一化为4值矩形供有线识别
                                if isinstance(ocr_poly[0], (list, tuple)) and len(ocr_poly) == 4:
                                    xs = [p[0] for p in ocr_poly]
                                    ys = [p[1] for p in ocr_poly]
                                    rect = [min(xs), min(ys), max(xs), max(ys)]
                                elif len(ocr_poly) == 8:
                                    xs = [ocr_poly[i] for i in range(0, 8, 2)]
                                    ys = [ocr_poly[i] for i in range(1, 8, 2)]
                                    rect = [min(xs), min(ys), max(xs), max(ys)]
                                elif len(ocr_poly) == 4:
                                    rect = [ocr_poly[0], ocr_poly[1], ocr_poly[2], ocr_poly[3]]
                                else:
                                    rect = None
                ocr_source = "cropped_ocr"
                logger.info(f"📊 OCR detected {len(ocr_boxes)} text boxes in table (cropped)")
            except Exception as e:
                logger.warning(f"Table OCR detection failed: {e}")
        
        return cropped_table, ocr_boxes, table_angle, ocr_source
    
    def process_table_element_wired(
        self,
        image: np.ndarray,
        layout_item: Dict[str, Any],
        scale: float,
        pre_matched_spans: Optional[List[Dict[str, Any]]] = None
    ) -> Dict[str, Any]:
        """
        使用 UNet 有线表格识别处理表格元素
        
        流程：
        1. OCR检测获取文本框坐标
        2. UNet 有线表格识别
        3. 坐标逆向转换回原图坐标
        
        Args:
            image: 页面图像
            layout_item: 布局检测项
            scale: 缩放比例
            pre_matched_spans: 预匹配的 OCR spans（来自整页 OCR）
            
        Returns:
            处理后的元素字典
        """
        bbox = layout_item.get('bbox', [0, 0, 0, 0])
        
        # OCR 预处理（返回已旋转的表格图片 + OCR 框）
        cropped_table, ocr_boxes, table_angle, ocr_source = \
            self._prepare_table_ocr(image, bbox, pre_matched_spans)
        
        # 获取裁剪后表格图片的尺寸
        orig_table_h, orig_table_w = cropped_table.shape[:2]
        orig_table_size = (orig_table_w, orig_table_h)
        
        # UNet 有线表格识别
        cells = []
        enhanced_html = ""
        
        try:
            if not self.wired_table_recognizer:
                raise RuntimeError("Wired table recognizer not available")
            
            wired_res = self.wired_table_recognizer.recognize(
                table_image=cropped_table,
                # ocr_boxes=ocr_boxes_for_wired,
                ocr_boxes=ocr_boxes,
            )
            
            if not (wired_res.get('html') or wired_res.get('cells')):
                raise RuntimeError("Wired recognizer returned empty result")
            
            cells = wired_res.get('cells', [])
            enhanced_html = wired_res.get('html', '')
            logger.info(f"📊 Wired UNet recognized {len(cells)} cells")
            
        except Exception as e:
            logger.warning(f"⚠️ Wired UNet recognition failed: {e}, falling back to empty result")
            # 有线识别失败，返回空表格（由 pipeline 层负责 fallback 到 VLM）
            return self._create_empty_table_result(layout_item, bbox, table_angle, ocr_source)
        
        # 坐标转换：将旋转后的坐标转换回原图坐标
        if table_angle != 0 and MERGER_AVAILABLE:
            cells, enhanced_html = CoordinateUtils.inverse_rotate_table_coords(
                cells=cells,
                html=enhanced_html,
                rotation_angle=table_angle,
                orig_table_size=orig_table_size,
                table_bbox=bbox
            )
            ocr_boxes = CoordinateUtils.inverse_rotate_ocr_boxes(
                ocr_boxes=ocr_boxes,
                rotation_angle=table_angle,
                orig_table_size=orig_table_size,
                table_bbox=bbox
            )
            logger.info(f"📐 Wired table coordinates transformed back to original image")
        else:
            # 没有旋转，只需要加上表格偏移量
            cells = CoordinateUtils.add_table_offset_to_cells(cells, bbox)
            enhanced_html = CoordinateUtils.add_table_offset_to_html(enhanced_html, bbox)
            ocr_boxes = CoordinateUtils.add_table_offset_to_ocr_boxes(ocr_boxes, bbox)
        
        return {
            'type': 'table',
            'bbox': bbox,
            'confidence': layout_item.get('confidence', 0.0),
            'content': {
                'html': enhanced_html,
                'original_html': enhanced_html,
                'cells': cells,
                'ocr_boxes': ocr_boxes,
                'table_angle': table_angle,
                'skew_angle': 0.0,
                'ocr_source': ocr_source,
                'recognition_method': 'wired_unet',
            },
        }
    
    def process_table_element_vlm(
        self,
        image: np.ndarray,
        layout_item: Dict[str, Any],
        scale: float,
        pre_matched_spans: Optional[List[Dict[str, Any]]] = None
    ) -> Dict[str, Any]:
        """
        使用 VLM 无线表格识别处理表格元素
        
        流程：
        1. OCR检测获取文本框坐标
        2. VLM识别获取表格结构HTML
        3. 匹配OCR坐标与VLM结构
        4. 坐标逆向转换回原图坐标
        
        Args:
            image: 页面图像
            layout_item: 布局检测项
            scale: 缩放比例
            pre_matched_spans: 预匹配的 OCR spans（来自整页 OCR）
            
        Returns:
            处理后的元素字典
        """
        bbox = layout_item.get('bbox', [0, 0, 0, 0])
        
        # OCR 预处理（返回已旋转的表格图片 + OCR 框）
        cropped_table, ocr_boxes, table_angle, ocr_source = \
            self._prepare_table_ocr(image, bbox, pre_matched_spans)
        
        # 获取裁剪后表格图片的尺寸
        orig_table_h, orig_table_w = cropped_table.shape[:2]
        orig_table_size = (orig_table_w, orig_table_h)
        
        # VLM 识别获取表格结构HTML
        table_html = ""
        try:
            vl_result = self.vl_recognizer.recognize_table(
                cropped_table,
                return_cells_coordinate=True
            )
            table_html = vl_result.get('html', '')
            logger.info(f"📊 VLM recognized table structure")
        except Exception as e:
            logger.warning(f"VLM table recognition failed: {e}")
            return self._create_empty_table_result(layout_item, bbox, table_angle, ocr_source)
        
        # 匹配OCR坐标与VLM结构
        cells = []
        enhanced_html = table_html
        skew_angle = 0.0
        
        if table_html and ocr_boxes and self.table_cell_matcher:
            try:
                enhanced_html, cells, _, skew_angle = self.table_cell_matcher.enhance_table_html_with_bbox(
                    html=table_html,
                    paddle_text_boxes=ocr_boxes,
                    start_pointer=0,
                    table_bbox=[0, 0, bbox[2] - bbox[0], bbox[3] - bbox[1]]
                )
                logger.info(f"📊 Matched {len(cells)} cells with coordinates (skew: {skew_angle:.2f}°)")
            except Exception as e:
                logger.warning(f"Cell coordinate matching failed: {e}")
        
        # 坐标转换：将旋转后的坐标转换回原图坐标
        if table_angle != 0 and MERGER_AVAILABLE:
            cells, enhanced_html = CoordinateUtils.inverse_rotate_table_coords(
                cells=cells,
                html=enhanced_html,
                rotation_angle=table_angle,
                orig_table_size=orig_table_size,
                table_bbox=bbox
            )
            ocr_boxes = CoordinateUtils.inverse_rotate_ocr_boxes(
                ocr_boxes=ocr_boxes,
                rotation_angle=table_angle,
                orig_table_size=orig_table_size,
                table_bbox=bbox
            )
            logger.info(f"📐 VLM table coordinates transformed back to original image")
        else:
            # 没有旋转，只需要加上表格偏移量
            cells = CoordinateUtils.add_table_offset_to_cells(cells, bbox)
            enhanced_html = CoordinateUtils.add_table_offset_to_html(enhanced_html, bbox)
            ocr_boxes = CoordinateUtils.add_table_offset_to_ocr_boxes(ocr_boxes, bbox)
        
        return {
            'type': 'table',
            'bbox': bbox,
            'confidence': layout_item.get('confidence', 0.0),
            'content': {
                'html': enhanced_html,
                'original_html': table_html,
                'cells': cells,
                'ocr_boxes': ocr_boxes,
                'table_angle': table_angle,
                'skew_angle': skew_angle,
                'ocr_source': ocr_source,
                'recognition_method': 'vlm',
            },
        }
    
    def _create_empty_table_result(
        self,
        layout_item: Dict[str, Any],
        bbox: List[float],
        table_angle: int = 0,
        ocr_source: str = "none"
    ) -> Dict[str, Any]:
        """
        创建空表格结果（识别失败时使用）
        
        Args:
            layout_item: 布局检测项
            bbox: 表格 bbox
            table_angle: 表格旋转角度
            ocr_source: OCR 来源
            
        Returns:
            空表格结果
        """
        return {
            'type': 'table',
            'bbox': bbox,
            'confidence': layout_item.get('confidence', 0.0),
            'content': {
                'html': '',
                'original_html': '',
                'cells': [],
                'ocr_boxes': [],
                'table_angle': table_angle,
                'skew_angle': 0.0,
                'ocr_source': ocr_source,
                'recognition_method': 'none',
            },
        }
    
    def process_equation_element(
        self,
        image: np.ndarray,
        layout_item: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        处理公式元素
        
        Args:
            image: 页面图像
            layout_item: 布局检测项
            
        Returns:
            处理后的元素字典
        """
        bbox = layout_item.get('bbox', [0, 0, 0, 0])
        category = layout_item.get('category', '')
        cropped_region = CoordinateUtils.crop_region(image, bbox)
        
        content = {'latex': '', 'confidence': 0.0}
        
        try:
            formula_result = self.vl_recognizer.recognize_formula(cropped_region)
            content = {
                'latex': formula_result.get('latex', ''),
                'confidence': formula_result.get('confidence', 0.0)
            }
        except Exception as e:
            logger.warning(f"Formula recognition failed: {e}")
        
        return {
            'type': category,
            'bbox': bbox,
            'confidence': layout_item.get('confidence', 0.0),
            'content': content
        }
    
    def process_image_element(
        self,
        image: np.ndarray,
        layout_item: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        处理图片元素
        
        Args:
            image: 页面图像
            layout_item: 布局检测项
            
        Returns:
            处理后的元素字典
        """
        bbox = layout_item.get('bbox', [0, 0, 0, 0])
        category = layout_item.get('category', '')
        
        cropped_image = CoordinateUtils.crop_region(image, bbox)
        
        return {
            'type': category,
            'bbox': bbox,
            'confidence': layout_item.get('confidence', 0.0),
            'content': {
                'type': 'image',
                'description': '',
                'image_data': cropped_image,
                'image_path': ''
            }
        }
    
    def process_code_element(
        self,
        image: np.ndarray,
        layout_item: Dict[str, Any],
        pdf_type: str,
        pdf_doc: Optional[Any],
        page_idx: int,
        scale: float,
        pre_matched_spans: Optional[List[Dict[str, Any]]] = None
    ) -> Dict[str, Any]:
        """
        处理代码元素
        
        处理优先级：
        1. 如果有预匹配的 spans（整页 OCR 结果），优先使用
        2. 数字PDF：尝试 PDF 字符提取
        3. 扫描件或提取失败：裁剪区域 OCR（兜底方案）
        
        Args:
            image: 页面图像
            layout_item: 布局检测项
            pdf_type: PDF类型
            pdf_doc: PDF文档对象
            page_idx: 页码索引
            scale: 缩放比例
            pre_matched_spans: 预匹配的 OCR spans（来自整页 OCR）
            
        Returns:
            处理后的元素字典
        """
        bbox = layout_item.get('bbox', [0, 0, 0, 0])
        category = layout_item.get('category', '')
        
        code_content = ""
        ocr_details = []
        extraction_method = "none"
        
        # 优先级1：使用预匹配的 spans（整页 OCR 结果）
        if pre_matched_spans and len(pre_matched_spans) > 0:
            code_content, sorted_spans = SpanMatcher.merge_spans_to_text(
                pre_matched_spans, bbox
            )
            if code_content.strip():
                ocr_details = sorted_spans
                extraction_method = "fullpage_ocr"
                logger.debug(f"📝 Code from full-page OCR: '{code_content[:30]}...'")
        
        # 优先级2：数字 PDF 字符提取
        if not code_content.strip() and pdf_type == 'txt' and pdf_doc is not None:
            code_content, success = PDFUtils.extract_text_from_pdf(
                pdf_doc, page_idx, bbox, scale
            )
            if success and code_content.strip():
                extraction_method = "pdf_extract"
            else:
                code_content = ""
        
        # 优先级3：裁剪区域 OCR（兜底方案）
        if not code_content.strip():
            try:
                cropped_image = CoordinateUtils.crop_region(image, bbox)
                ocr_results = self.ocr_recognizer.recognize_text(cropped_image)
                if ocr_results:
                    lines = []
                    for item in ocr_results:
                        lines.append(item.get('text', ''))
                    code_content = '\n'.join(lines)
                    # 将 OCR 坐标转换为绝对坐标
                    ocr_details = self._convert_ocr_details_to_absolute(ocr_results, bbox)
                    extraction_method = "cropped_ocr"
            except Exception as e:
                logger.warning(f"Code OCR failed: {e}")
        
        return {
            'type': category,
            'bbox': bbox,
            'confidence': layout_item.get('confidence', 0.0),
            'content': {
                'code': code_content,
                'language': '',
                'ocr_details': ocr_details,
                'extraction_method': extraction_method
            }
        }
    
    def process_discard_element(
        self,
        image: np.ndarray,
        layout_item: Dict[str, Any],
        pre_matched_spans: Optional[List[Dict[str, Any]]] = None
    ) -> Dict[str, Any]:
        """
        处理丢弃元素（水印、装饰等）
        
        提取文字块用于后续分析，但标记为丢弃类型
        
        Args:
            image: 页面图像
            layout_item: 布局检测项
            pre_matched_spans: 预匹配的 OCR spans（来自整页 OCR）
            
        Returns:
            处理后的元素字典
        """
        bbox = layout_item.get('bbox', [0, 0, 0, 0])
        category = layout_item.get('category', 'abandon')
        
        text_content = ""
        ocr_details: List[Dict[str, Any]] = []
        
        # 优先使用预匹配的 spans
        if pre_matched_spans and len(pre_matched_spans) > 0:
            text_content, sorted_spans = SpanMatcher.merge_spans_to_text(
                pre_matched_spans, bbox
            )
            if text_content.strip():
                ocr_details = sorted_spans
        
        # 兜底：裁剪区域 OCR
        if not text_content.strip():
            try:
                cropped_image = CoordinateUtils.crop_region(image, bbox)
                ocr_results = self.ocr_recognizer.recognize_text(cropped_image)
                if ocr_results:
                    text_parts = [
                        item['text'] for item in ocr_results 
                        if item.get('confidence', 0) > 0.5
                    ]
                    text_content = " ".join(text_parts)
                    # 将 OCR 坐标转换为绝对坐标
                    ocr_details = self._convert_ocr_details_to_absolute(ocr_results, bbox)
            except Exception as e:
                logger.debug(f"Discard element OCR failed: {e}")
        
        return {
            'type': 'discarded',
            'bbox': bbox,
            'confidence': layout_item.get('confidence', 0.0),
            'original_category': category,
            'content': {
                'text': text_content,
                'ocr_details': ocr_details
            }
        }
    
    @staticmethod
    def create_error_element(
        layout_item: Dict[str, Any], 
        error_msg: str
    ) -> Dict[str, Any]:
        """
        创建错误元素
        
        Args:
            layout_item: 原始布局项
            error_msg: 错误信息
            
        Returns:
            错误元素字典
        """
        return {
            'type': 'error',
            'bbox': layout_item.get('bbox', [0, 0, 0, 0]),
            'confidence': 0.0,
            'content': {'error': error_msg},
            'original_category': layout_item.get('category', 'unknown')
        }

