ocr_verify/merger/markdown_generator.py

"""
Markdown 生成模块
负责将合并后的数据生成 Markdown 文件
"""
import shutil
from pathlib import Path
from typing import List, Dict, Optional


class MarkdownGenerator:
    """Markdown 生成器"""
    
    @staticmethod
    def detect_data_format(merged_data: List[Dict]) -> str:
        """
        检测数据格式
        
        Returns:
            'mineru' 或 'paddleocr_vl'
        """
        if not merged_data:
            return 'mineru'
        
        first_item = merged_data[0]
        
        # 检查是否有 PaddleOCR_VL 特有字段
        if 'block_label' in first_item and 'block_content' in first_item:
            return 'paddleocr_vl'
        
        # 检查是否有 MinerU 特有字段
        if 'type' in first_item and ('table_body' in first_item or 'text' in first_item):
            return 'mineru'
        
        # 默认按 MinerU 格式处理
        return 'mineru'
    
    @staticmethod
    def generate_enhanced_markdown(merged_data: List[Dict], 
                                   output_path: Optional[str] = None,
                                   source_file: Optional[str] = None,
                                   data_format: Optional[str] = None) -> str:
        """
        生成增强的 Markdown（包含 bbox 信息的注释）
        
        Args:
            merged_data: 合并后的数据
            output_path: 输出路径
            source_file: 源文件路径（用于复制图片）
            data_format: 数据格式 ('mineru' 或 'paddleocr_vl')，None 则自动检测
        
        Returns:
            Markdown 内容
        """
        # ✅ 自动检测数据格式
        if data_format is None:
            data_format = MarkdownGenerator.detect_data_format(merged_data)
        
        print(f"ℹ️  检测到数据格式: {data_format}")
        
        # ✅ 根据格式选择处理函数
        if data_format == 'paddleocr_vl':
            return MarkdownGenerator._generate_paddleocr_vl_markdown(
                merged_data, output_path, source_file
            )
        else:
            return MarkdownGenerator._generate_mineru_markdown(
                merged_data, output_path, source_file
            )
    
    @staticmethod
    def _generate_mineru_markdown(merged_data: List[Dict],
                                  output_path: Optional[str] = None,
                                  source_file: Optional[str] = None) -> str:
        """生成 MinerU 格式的 Markdown"""
        md_lines = []
        
        for item in merged_data:
            item_type = item.get('type', '')
            
            if item_type == 'title':
                md_lines.extend(MarkdownGenerator._format_mineru_title(item))
            elif item_type == 'text':
                md_lines.extend(MarkdownGenerator._format_mineru_text(item))
            elif item_type == 'list':
                md_lines.extend(MarkdownGenerator._format_mineru_list(item))
            elif item_type == 'table':
                md_lines.extend(MarkdownGenerator._format_mineru_table(item))
            elif item_type == 'image':
                md_lines.extend(MarkdownGenerator._format_mineru_image(
                    item, output_path, source_file
                ))
            elif item_type in ['equation', 'interline_equation']:
                md_lines.extend(MarkdownGenerator._format_equation(item))
            elif item_type == 'inline_equation':
                md_lines.extend(MarkdownGenerator._format_inline_equation(item))
            elif item_type == 'header':
                md_lines.extend(MarkdownGenerator._format_mineru_header(item))
            elif item_type == 'footer':
                md_lines.extend(MarkdownGenerator._format_mineru_footer(item))
            elif item_type == 'page_number':
                md_lines.extend(MarkdownGenerator._format_mineru_page_number(item))
            elif item_type == 'ref_text':
                md_lines.extend(MarkdownGenerator._format_reference(item))
            else:
                md_lines.extend(MarkdownGenerator._format_unknown(item))
        
        markdown_content = '\n'.join(md_lines)
        
        if output_path:
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(markdown_content)
        
        return markdown_content
    
    @staticmethod
    def _generate_paddleocr_vl_markdown(merged_data: List[Dict],
                                        output_path: Optional[str] = None,
                                        source_file: Optional[str] = None) -> str:
        """生成 PaddleOCR_VL 格式的 Markdown"""
        md_lines = []
        
        for item in merged_data:
            block_label = item.get('block_label', '')
            
            if 'title' in block_label:
                md_lines.extend(MarkdownGenerator._format_paddleocr_vl_title(item))
            elif block_label == 'text':
                md_lines.extend(MarkdownGenerator._format_paddleocr_vl_text(item))
            elif block_label == 'table':
                md_lines.extend(MarkdownGenerator._format_paddleocr_vl_table(item))
            elif block_label == 'image':
                md_lines.extend(MarkdownGenerator._format_paddleocr_vl_figure(item))
            elif block_label == 'equation':
                md_lines.extend(MarkdownGenerator._format_paddleocr_vl_equation(item))
            elif block_label == 'reference':
                md_lines.extend(MarkdownGenerator._format_paddleocr_vl_reference(item))
            else:
                md_lines.extend(MarkdownGenerator._format_paddleocr_vl_unknown(item))
        
        markdown_content = '\n'.join(md_lines)
        
        if output_path:
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(markdown_content)
        
        return markdown_content
    
    # ================== MinerU 格式化方法 ==================
    
    @staticmethod
    def _format_mineru_title(item: Dict) -> List[str]:
        """格式化 MinerU 标题"""
        lines = []
        bbox = item.get('bbox', [])
        if bbox:
            lines.append(MarkdownGenerator._add_bbox_comment(bbox))
        
        text = item.get('text', '')
        text_level = item.get('text_level', 1)
        heading = '#' * min(text_level, 6)
        lines.append(f"{heading} {text}\n")
        
        return lines
    
    @staticmethod
    def _format_mineru_text(item: Dict) -> List[str]:
        """格式化 MinerU 文本"""
        lines = []
        bbox = item.get('bbox', [])
        if bbox:
            lines.append(MarkdownGenerator._add_bbox_comment(bbox))
        
        text = item.get('text', '')
        text_level = item.get('text_level', 0)
        
        if text_level > 0:
            heading = '#' * min(text_level, 6)
            lines.append(f"{heading} {text}\n")
        else:
            lines.append(f"{text}\n")
        
        return lines
    
    @staticmethod
    def _format_mineru_list(item: Dict) -> List[str]:
        """格式化 MinerU 列表"""
        lines = []
        bbox = item.get('bbox', [])
        if bbox:
            lines.append(MarkdownGenerator._add_bbox_comment(bbox))
        
        list_items = item.get('list_items', [])
        for list_item in list_items:
            lines.append(f"{list_item}\n")
        
        lines.append("")
        return lines
    
    @staticmethod
    def _format_mineru_table(item: Dict) -> List[str]:
        """格式化 MinerU 表格"""
        lines = []
        bbox = item.get('bbox', [])
        if bbox:
            lines.append(MarkdownGenerator._add_bbox_comment(bbox))
        
        # 表格标题
        table_caption = item.get('table_caption', [])
        for caption in table_caption:
            if caption:
                lines.append(f"**{caption}**\n")
        
        # 表格内容
        table_body = item.get('table_body_with_bbox', item.get('table_body', ''))
        if table_body:
            lines.append(table_body)
            lines.append("")
        
        # 表格脚注
        table_footnote = item.get('table_footnote', [])
        for footnote in table_footnote:
            if footnote:
                lines.append(f"*{footnote}*")
        if table_footnote:
            lines.append("")
        
        return lines
    
    @staticmethod
    def _format_mineru_image(item: Dict, output_path: Optional[str],
                            source_file: Optional[str]) -> List[str]:
        """格式化 MinerU 图片"""
        lines = []
        bbox = item.get('bbox', [])
        if bbox:
            lines.append(MarkdownGenerator._add_bbox_comment(bbox))
        
        img_path = item.get('img_path', '')
        
        # 复制图片
        if img_path and source_file and output_path:
            MarkdownGenerator._copy_image(img_path, source_file, output_path)
        
        # 图片标题
        image_caption = item.get('image_caption', [])
        for caption in image_caption:
            if caption:
                lines.append(f"**{caption}**\n")
        
        lines.append(f"![Image]({img_path})\n")
        
        # 图片脚注
        image_footnote = item.get('image_footnote', [])
        for footnote in image_footnote:
            if footnote:
                lines.append(f"*{footnote}*")
        if image_footnote:
            lines.append("")
        
        return lines
    
    @staticmethod
    def _format_mineru_header(item: Dict) -> List[str]:
        """格式化MinerU header"""
        lines = []
        bbox = item.get('bbox', [])
        if bbox:
            lines.append(MarkdownGenerator._add_bbox_comment(bbox))
        
        text = item.get('text', '')
        lines.append(f"<!-- 页眉: {text} -->\n")
        return lines

    @staticmethod
    def _format_mineru_footer(item: Dict) -> List[str]:
        """格式化MinerU footer"""
        lines = []
        bbox = item.get('bbox', [])
        if bbox:
            lines.append(MarkdownGenerator._add_bbox_comment(bbox))
        
        text = item.get('text', '')
        lines.append(f"<!-- 页脚: {text} -->\n")
        return lines

    @staticmethod
    def _format_mineru_page_number(item: Dict) -> List[str]:
        """格式化MinerU page_number"""
        lines = []
        bbox = item.get('bbox', [])
        if bbox:
            lines.append(MarkdownGenerator._add_bbox_comment(bbox))
        
        text = item.get('text', '')
        lines.append(f"<!-- 页码: {text} -->\n")
        return lines

    # ================== PaddleOCR_VL 格式化方法 ==================
    
    @staticmethod
    def _format_paddleocr_vl_title(item: Dict) -> List[str]:
        """格式化 PaddleOCR_VL 标题"""
        lines = []
        bbox = item.get('block_bbox', [])
        if bbox:
            lines.append(MarkdownGenerator._add_bbox_comment(bbox))
        
        text = item.get('block_content', '')
        block_label = item.get('block_label', '')
        
        # 根据 block_label 确定标题级别
        level_map = {
            'paragraph_title': 1,
            'figure_title': 2,
            'title': 1
        }
        text_level = level_map.get(block_label, 1)
        
        heading = '#' * min(text_level, 6)
        lines.append(f"{heading} {text}\n")
        
        return lines
    
    @staticmethod
    def _format_paddleocr_vl_text(item: Dict) -> List[str]:
        """格式化 PaddleOCR_VL 文本"""
        lines = []
        bbox = item.get('block_bbox', [])
        if bbox:
            lines.append(MarkdownGenerator._add_bbox_comment(bbox))
        
        text = item.get('block_content', '')
        lines.append(f"{text}\n")
        
        return lines
    
    @staticmethod
    def _format_paddleocr_vl_table(item: Dict) -> List[str]:
        """格式化 PaddleOCR_VL 表格"""
        lines = []
        bbox = item.get('block_bbox', [])
        if bbox:
            lines.append(MarkdownGenerator._add_bbox_comment(bbox))
        
        # 表格内容
        table_content = item.get('block_content_with_bbox', 
                                item.get('block_content', ''))
        if table_content:
            lines.append(table_content)
            lines.append("")
        
        return lines
    
    @staticmethod
    def _format_paddleocr_vl_figure(item: Dict) -> List[str]:
        """格式化 PaddleOCR_VL 图片"""
        lines = []
        bbox = item.get('block_bbox', [])
        if bbox:
            lines.append(MarkdownGenerator._add_bbox_comment(bbox))
        
        # PaddleOCR_VL 图片信息在 block_content 中
        content = item.get('block_content', '')
        lines.append(f"![Figure]({content})\n")
        
        return lines
    
    @staticmethod
    def _format_paddleocr_vl_equation(item: Dict) -> List[str]:
        """格式化 PaddleOCR_VL 公式"""
        lines = []
        bbox = item.get('block_bbox', [])
        if bbox:
            lines.append(MarkdownGenerator._add_bbox_comment(bbox))
        
        latex = item.get('block_content', '')
        if latex:
            lines.append(f"$$\n{latex}\n$$\n")
        
        return lines
    
    @staticmethod
    def _format_paddleocr_vl_reference(item: Dict) -> List[str]:
        """格式化 PaddleOCR_VL 参考文献"""
        text = item.get('block_content', '')
        return [f"> {text}\n"]
    
    @staticmethod
    def _format_paddleocr_vl_unknown(item: Dict) -> List[str]:
        """格式化 PaddleOCR_VL 未知类型"""
        lines = []
        bbox = item.get('block_bbox', [])
        if bbox:
            lines.append(MarkdownGenerator._add_bbox_comment(bbox))
        
        text = item.get('block_content', '')
        if text:
            lines.append(f"{text}\n")
        
        return lines
    
    # ================== 通用方法 ==================
    
    @staticmethod
    def _add_bbox_comment(bbox: List) -> str:
        """添加 bbox 注释"""
        return f"<!-- bbox: {bbox} -->"
    
    @staticmethod
    def _format_equation(item: Dict) -> List[str]:
        """格式化公式（通用）"""
        latex = item.get('latex', '')
        if latex:
            return [f"$$\n{latex}\n$$\n"]
        return []
    
    @staticmethod
    def _format_inline_equation(item: Dict) -> List[str]:
        """格式化行内公式（通用）"""
        latex = item.get('latex', '')
        if latex:
            return [f"${latex}$\n"]
        return []
    
    @staticmethod
    def _format_metadata(item: Dict, item_type: str) -> List[str]:
        """格式化元数据（通用）"""
        text = item.get('text', '')
        type_map = {
            'page_number': '页码',
            'header': '页眉',
            'footer': '页脚'
        }
        if text:
            return [f"<!-- {type_map.get(item_type, item_type)}: {text} -->\n"]
        return []
    
    @staticmethod
    def _format_reference(item: Dict) -> List[str]:
        """格式化参考文献（MinerU）"""
        text = item.get('text', '')
        return [f"> {text}\n"]
    
    @staticmethod
    def _format_unknown(item: Dict) -> List[str]:
        """格式化未知类型（MinerU）"""
        text = item.get('text', '')
        if text:
            return [f"{text}\n"]
        return []
    
    @staticmethod
    def _copy_image(img_path: str, source_file: str, output_path: str):
        """复制图片到输出目录"""
        source_dir = Path(source_file).parent
        img_full_path = source_dir / img_path
        if img_full_path.exists():
            output_img_path = Path(output_path).parent / img_path
            output_img_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy(img_full_path, output_img_path)