proposa-ai/pdf_table_to_docx/extractor.py

"""
extractor.py - PDF 表格提取与 DOCX 生成主模块

负责：
- 打开 PDF，逐页提取表格
- 检测跨页表格并自动合并
- 协调 table_parser 和 docx_writer 完成完整流程
- 支持单文件输出和目录式输出（JSON + 每表独立 DOCX）
"""

import json
import logging
import os
import re

import pdfplumber

from .table_parser import TableInfo, tables_on_page
from .docx_writer import build_document, build_single_table_document
from .text_scorer import get_scorer

logger = logging.getLogger(__name__)


# 跨页表格检测用容差
WIDTH_MATCH_EPSILON = 0.1   # 列宽比例差异小于此值认为匹配
PAGE_EDGE_EPSILON = 80      # pt，距页边距小于此值认为"靠近边缘"


def _table_first_row_preview(table: TableInfo, max_chars: int = 20) -> list:
    """提取表格首行前几列作为预览。"""
    preview = []
    if table.rows == 0:
        return preview
    for c in range(min(table.cols, 5)):
        cell = table.cells[0][c]
        if cell:
            txt = cell.text.strip().replace("\n", " ")[:max_chars]
            preview.append(txt)
        else:
            preview.append("")
    return preview


def _has_continuity_signal(t1: TableInfo, t2: TableInfo) -> bool:
    """检查跨页表格之间的内容连续性信号。

    当列宽不完全匹配时，用以下信号判断是否为同表跨页：

    信号 A——空首列延续：t2 的首行第 0 列为空，但 t1 对应位置有内容
    （考虑纵向合并，从下往上取 t1 末列第一个非空单元格）。

    信号 B——NLP 文本延续：使用语言模型评分，判断 t1 与 t2 交界处
    的文本是否可以拼接成完整语义单元（如「资源投入」+「和管理」）。

    信号 C——全表空首列：t2 在首列几乎全部为空，是 t1 分类标签的续表。
    """
    if t1.rows == 0 or t2.rows == 0:
        return False

    signals = []

    # ---- 辅助：获取 t1 末行指定列的第一个非空文本 ----
    def _last_nonempty(tbl: TableInfo, col: int) -> str:
        for r in range(tbl.rows - 1, -1, -1):
            txt = tbl.cells[r][col].text.strip()
            if txt:
                return txt
        return ""

    # ---- 信号 A：空首列延续 ----
    last_nonempty_t1_col0 = _last_nonempty(t1, 0)
    first_cell_t2 = t2.cells[0][0].text.strip()
    if not first_cell_t2 and last_nonempty_t1_col0:
        signals.append("A")

    # ---- 信号 B：NLP 文本延续（基于语言模型评分） ----
    try:
        scorer = get_scorer()
        for c in range(min(t1.cols, t2.cols)):
            t1_last = _last_nonempty(t1, c)
            t2_first = t2.cells[0][c].text.strip()
            if not t1_last or not t2_first:
                continue
            if len(t1_last) < 2 or len(t2_first) < 2:
                continue
            if scorer.is_continuation(t1_last, t2_first, threshold=1.5):
                logger.debug(f"NLP 文本延续信号: '{t1_last[-20:]}' + '{t2_first[:20]}'")
                signals.append("B")
                break
    except FileNotFoundError as e:
        logger.warning(f"语言模型未加载，跳过 NLP 检测: {e}")
    except Exception as e:
        logger.debug(f"NLP 评分异常（不影响主流程）: {e}")

    # ---- 信号 C：全表空首列 -> 可能是续表 ----
    all_first_col_empty = all(
        not t2.cells[r][0].text.strip()
        for r in range(min(t2.rows, 5))
    )
    if all_first_col_empty and last_nonempty_t1_col0:
        signals.append("C")

    return len(signals) > 0


def _tables_span_pages(t1: TableInfo, t2: TableInfo,
                       page1_height: float = None,
                       page2_height: float = None) -> bool:
    """判断两个表格是否属于跨页的同一个表格。

    两阶段判断：
    阶段一（严格匹配）：列数相同 + 列宽匹配 + 页面位置 + 非重复表头
    阶段二（宽松匹配）：列数相同 + 页面位置 + 内容连续性信号

    宽松匹配用于处理 pdfplumber 列宽解析不一致的情况
    （如跨页后因合并单元格导致列边界不同）。
    """
    if t1.cols != t2.cols:
        return False

    # ---- 公共前提：页面位置靠近 ----
    if t1.bbox and page1_height:
        if page1_height - t1.bbox[3] > PAGE_EDGE_EPSILON:
            return False
    if t2.bbox and page2_height:
        if t2.bbox[1] > PAGE_EDGE_EPSILON:
            return False

    # ---- 公共前提：首行不是重复表头 ----
    if t1.rows > 0 and t2.rows > 0:
        if _table_first_row_preview(t1) == _table_first_row_preview(t2):
            return False

    # ---- 阶段一：列宽严格匹配 ----
    if len(t1.col_widths) == len(t2.col_widths) and len(t1.col_widths) > 0:
        widths_match = True
        for i in range(len(t1.col_widths)):
            if abs(t1.col_widths[i] - t2.col_widths[i]) > WIDTH_MATCH_EPSILON:
                widths_match = False
                break
        if widths_match:
            return True

    # ---- 阶段二：内容连续性宽松匹配 ----
    if _has_continuity_signal(t1, t2):
        logger.debug(
            f"列宽不匹配但检测到延续信号，视为跨页："
            f"第 {t1.page_num+1} 页 → 第 {t2.page_num+1} 页"
        )
        return True

    return False


def _merge_tables(t1: TableInfo, t2: TableInfo) -> TableInfo:
    """合并两个跨页表格，并修复纵向合并的语义断层。"""
    skip_header = False
    if t1.rows > 0 and t2.rows > 0:
        header_same = True
        for c in range(min(t1.cols, t2.cols)):
            if t1.cells[0][c].text.strip() != t2.cells[0][c].text.strip():
                header_same = False
                break
        if header_same:
            skip_header = True

    t2_start = 1 if skip_header else 0
    merged_rows = t1.rows + (t2.rows - t2_start)
    merged_cells = []
    for r in range(t1.rows):
        merged_cells.append(list(t1.cells[r]))
    for r in range(t2_start, t2.rows):
        merged_cells.append(list(t2.cells[r]))

    # ---- 修复纵向合并的语义断层 ----
    # t1 的纵向合并标签（如"分项服务方案"）被分页截断，
    # t2 对应列首行为空。此时应扩展 t1 的合并覆盖到 t2 区域。
    for c in range(t1.cols):
        # t1 部分：从下往上找最后一个非空文本的合并起始单元格
        t1_origin_r = None
        t1_span = 0
        for r in range(t1.rows - 1, -1, -1):
            cell = merged_cells[r][c]
            if cell.row_span > 0 and cell.text.strip():
                t1_origin_r = r
                t1_span = cell.row_span
                break

        # t2 部分：第一个合并起始单元格
        t2_start_r = None
        t2_span = 0
        for r in range(t1.rows, merged_rows):
            cell = merged_cells[r][c]
            if cell.row_span > 0:
                t2_start_r = r
                t2_span = max(cell.row_span, 1)
                break

        if (t1_origin_r is not None and t2_start_r is not None
                and t1_span > 0 and t2_span > 0):
            t1_merge_end = t1_origin_r + t1_span - 1
            reaches_bottom = (t1_merge_end >= t1.rows - 1)
            t2_text = merged_cells[t2_start_r][c].text.strip()

            if not t2_text and reaches_bottom:
                new_end = t2_start_r + t2_span - 1
                new_span = new_end - t1_origin_r + 1
                merged_cells[t1_origin_r][c].row_span = new_span
                for rr in range(t2_start_r, min(t2_start_r + t2_span, merged_rows)):
                    merged_cells[rr][c].row_span = 0
                    merged_cells[rr][c].col_span = 0

    all_pages = list(t1.original_pages or [t1.page_num + 1])
    all_pages.extend(t2.original_pages or [t2.page_num + 1])

    return TableInfo(
        rows=merged_rows,
        cols=t1.cols,
        cells=merged_cells,
        col_widths=t1.col_widths,
        bbox=None,
        page_num=t1.page_num,
        cross_page=True,
        original_pages=sorted(set(all_pages)),
    )


def extract_tables_from_pdf(pdf_path: str) -> list:
    """从 PDF 中提取所有表格，包括跨页表格的合并。

    Returns:
        list[TableInfo]: 所有合并后的完整表格列表
    """
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"PDF 文件不存在: {pdf_path}")

    all_tables = []

    with pdfplumber.open(pdf_path) as pdf:
        total_pages = len(pdf.pages)
        logger.info(f"PDF 共 {total_pages} 页，开始提取表格...")

        # ---- 阶段 1：逐页提取表格 ----
        page_tables = []
        for page_num, page in enumerate(pdf.pages):
            tables = tables_on_page(page, page_num)
            if tables:
                page_tables.append((page_num, tables))
                logger.info(f"第 {page_num+1} 页: 发现 {len(tables)} 个表格")

        if not page_tables:
            logger.info("未检测到任何表格")
            return []

        # ---- 阶段 2：展平并合并跨页表格 ----
        flat_tables = []
        for page_num, tables in page_tables:
            for table in tables:
                table.page_num = page_num
                table.original_pages = [page_num + 1]
                flat_tables.append(table)

        merged = []
        i = 0
        while i < len(flat_tables):
            current = flat_tables[i]

            if i + 1 < len(flat_tables):
                next_table = flat_tables[i + 1]
                if next_table.page_num == current.page_num + 1:
                    # 获取两页的高度用于位置判断
                    p1_h = pdf.pages[current.page_num].height
                    p2_h = pdf.pages[next_table.page_num].height
                    if _tables_span_pages(
                        current, next_table,
                        page1_height=p1_h, page2_height=p2_h,
                    ):
                        logger.info(
                            f"检测到跨页表格：第 {current.page_num+1} 页 → "
                            f"第 {next_table.page_num+1} 页"
                        )
                        merged_table = _merge_tables(current, next_table)
                        merged.append(merged_table)
                        i += 2
                        continue

            merged.append(current)
            i += 1

        all_tables = merged

    logger.info(f"表格提取完成，共 {len(all_tables)} 个完整表格")
    for idx, t in enumerate(all_tables):
        pages_str = (
            f"（跨页: {t.original_pages}）" if t.cross_page else ""
        )
        logger.info(f"  表格 {idx+1}: {t.rows} 行 x {t.cols} 列 {pages_str}")

    return all_tables


def _sanitize_filename(name: str) -> str:
    """清理文件名中的非法字符。"""
    return re.sub(r'[\\/:*?"<>|]', "_", name)


def _build_json_metadata(tables: list, pdf_path: str, output_dir: str) -> dict:
    """构建表格元数据 JSON。"""
    pdf_name = os.path.basename(pdf_path)
    metadata = {
        "pdf_file": pdf_name,
        "pdf_path": os.path.abspath(pdf_path),
        "total_tables": len(tables),
        "tables": [],
    }

    for idx, t in enumerate(tables):
        preview = _table_first_row_preview(t)
        entry = {
            "id": idx + 1,
            "file": f"tables/table_{idx+1:03d}.docx",
            "pages": t.original_pages or [t.page_num + 1],
            "rows": t.rows,
            "cols": t.cols,
            "cross_page": t.cross_page,
            "preview": preview,
        }
        metadata["tables"].append(entry)

    return metadata


class PDFTableExtractor:
    """PDF 表格提取器，封装完整流程。"""

    def __init__(self, pdf_path: str):
        self.pdf_path = pdf_path
        self.tables = []

    def extract(self) -> list:
        """提取 PDF 中的表格。"""
        self.tables = extract_tables_from_pdf(self.pdf_path)
        return self.tables

    def to_docx(
        self,
        output_path: str = None,
        page_width_cm: float = 21.0,
        page_height_cm: float = 29.7,
        margin_cm: float = 2.0,
    ) -> str:
        """提取表格并生成单个 DOCX 文件（所有表格依次排列）。

        Returns:
            str: DOCX 文件路径
        """
        if not self.tables:
            self.extract()
        if not self.tables:
            logger.warning("未提取到任何表格，无法生成 DOCX")
            return ""

        if not output_path:
            base = os.path.splitext(self.pdf_path)[0]
            output_path = f"{base}_表格输出.docx"

        return build_document(
            self.tables,
            output_path,
            page_width_cm=page_width_cm,
            page_height_cm=page_height_cm,
            margin_cm=margin_cm,
        )

    def to_docx_directory(
        self,
        output_dir: str = None,
        page_width_cm: float = 21.0,
        page_height_cm: float = 29.7,
        margin_cm: float = 2.0,
    ) -> str:
        """提取表格并生成目录式输出：
        - output_dir/tables.json         全局元数据
        - output_dir/tables/table_NNN.docx  每个表格独立文件

        Args:
            output_dir: 输出目录，默认与 PDF 同目录 + "_表格输出"
            page_width_cm: 页面宽度
            page_height_cm: 页面高度
            margin_cm: 页边距

        Returns:
            str: 输出目录路径
        """
        if not self.tables:
            self.extract()
        if not self.tables:
            logger.warning("未提取到任何表格，无法生成输出目录")
            return ""

        if not output_dir:
            base = os.path.splitext(self.pdf_path)[0]
            output_dir = f"{base}_表格输出"

        os.makedirs(output_dir, exist_ok=True)

        # ---- 创建表格子目录 ----
        tables_dir = os.path.join(output_dir, "tables")
        os.makedirs(tables_dir, exist_ok=True)

        # ---- 写入每个表格的独立 DOCX ----
        for idx, table_info in enumerate(self.tables):
            if table_info.rows == 0 or table_info.cols == 0:
                continue

            docx_name = f"table_{idx+1:03d}.docx"
            docx_path = os.path.join(tables_dir, docx_name)

            label = f"表格 {idx+1}"
            pages = table_info.original_pages or [table_info.page_num + 1]
            if table_info.cross_page:
                label += f"（跨页: 第 {pages[0]}-{pages[-1]} 页）"
            else:
                label += f"（第 {pages[0]} 页）"

            build_single_table_document(
                table_info,
                docx_path,
                page_width_cm=page_width_cm,
                page_height_cm=page_height_cm,
                margin_cm=margin_cm,
                table_label=label,
            )

        # ---- 写入 JSON 元数据 ----
        meta = _build_json_metadata(self.tables, self.pdf_path, output_dir)
        json_path = os.path.join(output_dir, "tables.json")
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump(meta, f, ensure_ascii=False, indent=2)

        logger.info(f"目录输出完成，共 {len(self.tables)} 个表格")
        logger.info(f"  元数据: {json_path}")
        logger.info(f"  表格目录: {tables_dir}/")

        return output_dir