"""
OCR 数据解析模块
支持 PPStructure 和 MinerU 两种格式
"""

from typing import Dict, List, Tuple, Optional, Union
from bs4 import BeautifulSoup


class OcrParser:
    """OCR 数据解析器"""
    
    @staticmethod
    def parse(ocr_result: Dict, json_format: str = "auto") -> Tuple[List[int], Dict]:
        """
        统一的 OCR 数据解析接口
        
        Args:
            ocr_result: OCR 识别结果（完整 JSON）
            json_format: 数据格式类型 ("auto" / "ppstructure" / "mineru")
        
        Returns:
            (table_bbox, ocr_data): 表格边界框和统一格式的数据
        """
        if json_format.lower() == "auto":
            json_format_type = OcrParser.detect_format(ocr_result)
        else:
            json_format_type = json_format.lower()
        
        if json_format_type == "structured":
            return OcrParser._parse_structured(ocr_result)
        elif json_format_type == "ppstructure":
            return OcrParser._parse_ppstructure(ocr_result)
        elif json_format_type == "mineru":
            return OcrParser._parse_mineru(ocr_result)
        else:
            raise ValueError(f"不支持的数据格式类型: {json_format}")
    
    @staticmethod
    def detect_format(ocr_result: Dict) -> str:
        """检测 OCR 数据格式类型"""
        # 检测是否为结构化格式（包含 horizontal_lines、vertical_lines 等）
        if 'horizontal_lines' in ocr_result and 'vertical_lines' in ocr_result:
            return "structured"
        elif 'parsing_res_list' in ocr_result and 'overall_ocr_res' in ocr_result:
            return "ppstructure"
        else:
            return "mineru"

    @staticmethod
    def _parse_structured(structured_data: Dict) -> Tuple[List[int], Dict]:
        """
        解析结构化格式数据（包含 horizontal_lines、vertical_lines 等）
        
        Args:
            structured_data: 结构化的表格数据
        
        Returns:
            (table_bbox, ocr_data): 表格边界框和统一格式的数据
        """
        # 获取必要的字段
        if 'table_bbox' not in structured_data:
            raise ValueError("结构化数据中缺少 table_bbox 字段")
        
        table_bbox = structured_data['table_bbox']
        horizontal_lines = structured_data.get('horizontal_lines', [])
        vertical_lines = structured_data.get('vertical_lines', [])
        total_rows = structured_data.get('total_rows', len(horizontal_lines) - 1)
        col_widths = structured_data.get('col_widths', [])
        
        ocr_data = {
            'table_bbox': table_bbox,
            'actual_rows': total_rows,
            'actual_cols': len(col_widths) if col_widths else len(vertical_lines) - 1,
            'text_boxes': [],  # 结构化格式中没有文本框信息
            'horizontal_lines': horizontal_lines,
            'vertical_lines': vertical_lines,
            'row_height': structured_data.get('row_height'),
            'col_widths': col_widths,
            'image_rotation_angle': 0.0,
            'skew_angle': 0.0,
            'original_skew_angle': 0.0
        }
        
        return table_bbox, ocr_data

    @staticmethod
    def _parse_mineru(mineru_result: Union[Dict, List]) -> Tuple[List[int], Dict]:
        """
        解析 MinerU 格式数据
        
        Args:
            mineru_result: MinerU 的完整 JSON 结果
            
        Returns:
            (table_bbox, ocr_data): 表格边界框和统一格式的数据
        """
        # 提取 table 数据
        table_data = OcrParser._extract_table_data(mineru_result)
        
        if not table_data:
            raise ValueError("未找到 MinerU 格式的表格数据 (type='table')")
        
        if 'table_cells' not in table_data:
            raise ValueError("表格数据中未找到 table_cells 字段")
        
        table_cells = table_data['table_cells']
        if not table_cells:
            raise ValueError("table_cells 为空")
        
        # 优先使用 table_body 确定准确的行列数
        if 'table_body' in table_data:
            actual_rows, actual_cols = OcrParser._parse_table_body_structure(table_data['table_body'])
        else:
            # 回退：从 table_cells 推断
            actual_rows = max((cell.get('row', 0) for cell in table_cells if 'row' in cell), default=0)
            actual_cols = max((cell.get('col', 0) for cell in table_cells if 'col' in cell), default=0)
        
        # 计算表格边界框
        if 'bbox' in table_data:
            table_bbox = table_data['bbox']
        else:
            all_bboxes = [cell['bbox'] for cell in table_cells if 'bbox' in cell]
            if all_bboxes:
                x_min = min(bbox[0] for bbox in all_bboxes)
                y_min = min(bbox[1] for bbox in all_bboxes)
                x_max = max(bbox[2] for bbox in all_bboxes)
                y_max = max(bbox[3] for bbox in all_bboxes)
                table_bbox = [x_min, y_min, x_max, y_max]
            else:
                raise ValueError(f"无法计算表格边界框，缺少 bbox 信息, table_cells 内容: {table_cells}")
        
        # 按位置排序
        table_cells.sort(key=lambda x: (x['bbox'][1], x['bbox'][0]))
        
        # 转换为统一格式
        ocr_data = {
            'table_bbox': table_bbox,
            'actual_rows': actual_rows,
            'actual_cols': actual_cols,
            'text_boxes': table_cells,
            'image_rotation_angle': table_data.get('image_rotation_angle', 0.0),
            'skew_angle': table_data.get('skew_angle', 0.0),
            'original_skew_angle': table_data.get('skew_angle', 0.0)
        }
        
        return table_bbox, ocr_data

    @staticmethod
    def _parse_ppstructure(ocr_result: Dict) -> Tuple[List[int], Dict]:
        """
        解析 PPStructure V3 格式数据
        
        Args:
            ocr_result: PPStructure V3 的完整 JSON 结果
        
        Returns:
            (table_bbox, ocr_data): 表格边界框和统一格式的数据
        """
        # 从 parsing_res_list 中找到 table 区域
        table_bbox = None
        if 'parsing_res_list' in ocr_result:
            for block in ocr_result['parsing_res_list']:
                if block.get('block_label') == 'table':
                    table_bbox = block.get('block_bbox')
                    break
        
        if not table_bbox:
            raise ValueError("未找到表格区域 (block_label='table')")
        
        # 从 overall_ocr_res 中提取文本框
        text_boxes = []
        if 'overall_ocr_res' in ocr_result:
            rec_boxes = ocr_result['overall_ocr_res'].get('rec_boxes', [])
            rec_texts = ocr_result['overall_ocr_res'].get('rec_texts', [])
            
            # 过滤出表格区域内的文本框
            for i, bbox in enumerate(rec_boxes):
                if len(bbox) >= 4:
                    x1, y1, x2, y2 = bbox[:4]
                    
                    # 判断文本框是否在表格区域内
                    if (x1 >= table_bbox[0] and y1 >= table_bbox[1] and
                        x2 <= table_bbox[2] and y2 <= table_bbox[3]):
                        text_boxes.append({
                            'bbox': [int(x1), int(y1), int(x2), int(y2)],
                            'text': rec_texts[i] if i < len(rec_texts) else ''
                        })
        
        # 按位置排序
        text_boxes.sort(key=lambda x: (x['bbox'][1], x['bbox'][0]))
        
        ocr_data = {
            'table_bbox': table_bbox,
            'text_boxes': text_boxes
        }
        
        return table_bbox, ocr_data

    @staticmethod
    def _extract_table_data(mineru_result: Union[Dict, List]) -> Optional[Dict]:
        """从 MinerU 结果中提取 table 数据"""
        if isinstance(mineru_result, list):
            for item in mineru_result:
                if isinstance(item, dict) and item.get('type') == 'table':
                    return item
        elif isinstance(mineru_result, dict):
            if mineru_result.get('type') == 'table':
                return mineru_result
            # 递归查找
            for value in mineru_result.values():
                if isinstance(value, dict) and value.get('type') == 'table':
                    return value
                elif isinstance(value, list):
                    result = OcrParser._extract_table_data(value)
                    if result:
                        return result
        return None

    @staticmethod
    def _parse_table_body_structure(table_body: str) -> Tuple[int, int]:
        """从 table_body HTML 中解析准确的行列数"""
        try:
            soup = BeautifulSoup(table_body, 'html.parser')
            table = soup.find('table')
            
            if not table:
                return 0, 0
            
            rows = table.find_all('tr')
            if not rows:
                return 0, 0
            
            num_rows = len(rows)
            num_cols = max(len(row.find_all(['td', 'th'])) for row in rows)
            
            return num_rows, num_cols
            
        except Exception:
            return 0, 0
