utils.py

"""
工具函数模块
============
通用辅助函数，包括文件处理、图像处理、文本预处理等。
"""

import base64
import os
import re
from pathlib import Path
from typing import Optional


# ======================================================================
#  文件和图像处理
# ======================================================================

def image_to_base64(image_path: str) -> str:
    """将图片转换为 base64 编码字符串"""
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def get_image_mime(image_path: str) -> str:
    """根据文件扩展名获取 MIME 类型"""
    ext = Path(image_path).suffix.lower()
    mime_map = {
        ".png": "image/png",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".gif": "image/gif",
        ".bmp": "image/bmp",
        ".webp": "image/webp",
    }
    return mime_map.get(ext, "image/png")


def get_output_path(input_path: str, output_path: Optional[str] = None, 
                   suffix: str = '.docx') -> str:
    """生成输出文件路径"""
    if output_path:
        return output_path
    return str(Path(input_path).with_suffix(suffix))


# ======================================================================
#  文本预处理
# ======================================================================

def normalize_math_delimiters(text: str) -> str:
    """
    统一数学公式定界符为 $ 和 $$
    \\(...\\) → $...$
    \\[...\\] → $$...$$
    """
    # \\( ... \\) → $ ... $
    text = re.sub(r'\\\((.+?)\\\)', r'$\1$', text, flags=re.DOTALL)
    # \\[ ... \\] → $$ ... $$
    text = re.sub(r'\\\[(.+?)\\\]', r'$$\1$$', text, flags=re.DOTALL)
    return text


def extract_markdown_body(text: str) -> str:
    """
    提取 Markdown 主体内容。
    如果文本前后有说明性文字（非 Markdown 正文），则提取 --- 之间的主体。
    如果内容本身以 # 标题开头，说明 --- 是章节分隔线，保留全文。
    """
    stripped = text.strip()
    
    # 如果全文以 # 标题开头，说明没有包裹说明文字，直接返回
    if stripped.startswith('#'):
        return text
    
    lines = text.split('\n')
    separators = [i for i, line in enumerate(lines)
                  if re.match(r'^-{3,}\s*$', line.strip())]
    
    if len(separators) >= 2:
        before_first = '\n'.join(lines[:separators[0]]).strip()
        if before_first and not before_first.startswith('#'):
            start = separators[0] + 1
            end = separators[-1]
            body = '\n'.join(lines[start:end])
            if body.strip():
                return body
    
    return text


def postprocess_markdown(text: str) -> str:
    """Markdown 专用后处理：统一定界符、移除分隔线、压缩空行"""
    text = normalize_math_delimiters(text)
    
    # 移除 Markdown 行末双空格换行标记
    text = re.sub(r'  +$', '', text, flags=re.MULTILINE)
    
    # 移除残留的 --- 分隔线
    text = re.sub(r'^-{3,}\s*$', '', text, flags=re.MULTILINE)
    
    # 压缩连续空行
    text = re.sub(r'\n{3,}', '\n\n', text)
    
    # display math $$ 独占行
    text = re.sub(r'([^\n])\$\$', r'\1\n$$', text)
    text = re.sub(r'\$\$([^\n])', r'$$\n\1', text)
    
    return text.strip()


def postprocess_ocr_markdown(markdown_text: str) -> str:
    """OCR 识别结果后处理"""
    # 清理常见 OCR 乱码
    markdown_text = re.sub(r'[\x00-\x08\x0b-\x0c\x0e-\x1f]', '', markdown_text)
    
    # 统一数学定界符
    markdown_text = normalize_math_delimiters(markdown_text)
    
    # 清理多余空行
    markdown_text = re.sub(r'\n{4,}', '\n\n\n', markdown_text)
    
    # 确保 $$ 独占行
    markdown_text = re.sub(r'([^\n])\$\$', r'\1\n$$', markdown_text)
    markdown_text = re.sub(r'\$\$([^\n])', r'$$\n\1', markdown_text)
    
    return markdown_text.strip()


# ======================================================================
#  文件读取
# ======================================================================

def read_text_file(file_path: str) -> str:
    """读取文本文件，自动处理编码"""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return f.read()
    except UnicodeDecodeError:
        with open(file_path, 'r', encoding='gbk') as f:
            return f.read()