Pdf Processor

# 外部 API 调用顺序（auto 模式）
# 该顺序就是 API 尝试顺序；支持值：mineru, paddle（逗号分隔）
OCR_API_ORDER="paddle,mineru"

# MinerU API（可选）
MINERU_API_BASE="https://mineru.net/api/v4"
MINERU_API_TOKEN="your_mineru_token_here"
# 可选：部分网关需要额外 user token 头
MINERU_USER_TOKEN=""
# 注意：MinerU API Token 约 90 天（3 个月）有效，过期后需重新申请
# 申请/更新地址：https://mineru.net/apiManage/token

# 外部 PaddleOCR API（可选）
PADDLE_OCR_API_ENDPOINT="https://your-aistudio-app.com/ocr"
PADDLE_OCR_API_KEY="your_paddle_token_here"

# 兼容别名（可选）：如果你已经使用旧变量，可继续保留
# API_URL="https://your-aistudio-app.com/ocr"
# TOKEN="your_token_here"
# MINERU_API_BASE_URL="https://mineru.net"
# MINERU_BASE_URL="https://mineru.net"
# MINERU_TOKEN="your_mineru_token_here"

MIT License

Copyright (c) 2026 maoking

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

{
  "description": "OCR 文本上下文纠错规则。每条规则包含正则表达式和替换文本，按顺序依次应用到 OCR 结果上。",
  "usage": "编辑此文件即可扩充规则；也可通过 --ocr-corrections <file> 加载自定义规则（与基础规则合并）",
  "rules": [
    {"pattern": "性别[芝兰芳芙]", "replacement": "性别女", "note": "性别女常见误识别"},
    {"pattern": "性别[暮幕另罗量]", "replacement": "性别男", "note": "性别男常见误识别"},
    {"pattern": "年能", "replacement": "年龄", "note": "年龄误识别"},
    {"pattern": "(\\d)些(?=[，。、\\s\\n）)]|$)", "replacement": "\\1岁", "note": "数字后'岁'误为'些'"},
    {"pattern": "脑疵", "replacement": "脑疝", "note": "脑疝误识别"},
    {"pattern": "狭實", "replacement": "狭窄", "note": "狭窄误识别"},
    {"pattern": "(颅内|颈内|基底|冠状|椎)动脑", "replacement": "\\1动脉", "note": "动脉误为动脑"},
    {"pattern": "烹治", "replacement": "意识", "note": "意识误识别"},
    {"pattern": "降碍", "replacement": "障碍", "note": "障碍误识别"},
    {"pattern": "偏[備癱]", "replacement": "偏瘫", "note": "偏瘫误识别"},
    {"pattern": "[癲癫][癇癎痫]", "replacement": "癫痫", "note": "癫痫变体"},
    {"pattern": "大(面积?)脆板塞", "replacement": "大\\1脑梗塞", "note": "脑梗塞误识别"},
    {"pattern": "气產", "replacement": "气管", "note": "气管误识别"},
    {"pattern": "碱遐", "replacement": "减弱", "note": "体格检查-减弱"},
    {"pattern": "滅退", "replacement": "减退", "note": "体格检查-减退"},
    {"pattern": "过礅", "replacement": "过敏", "note": "过敏误识别"},
    {"pattern": "法阮", "replacement": "法院", "note": "法院误识别"},
    {"pattern": "叛决", "replacement": "判决", "note": "判决误识别"},
    {"pattern": "诉论", "replacement": "诉讼", "note": "诉讼误识别"},
    {"pattern": "答告人", "replacement": "被告人", "note": "被告人误识别"},
    {"pattern": "证剧(?=[，。、\\s\\n）)])", "replacement": "证据", "note": "证据误识别"},
    {"pattern": "沙库巴去", "replacement": "沙库巴曲", "note": "药名：曲误为去"},
    {"pattern": "在乙拉西坦", "replacement": "左乙拉西坦", "note": "药名：左误为在"},
    {"pattern": "气室切开", "replacement": "气管切开", "note": "气管误识别"},
    {"pattern": "湍謩捾醌", "replacement": "温馨提示", "note": "温馨提示乱码"},
    {"pattern": "继续继续", "replacement": "继续", "note": "重复词"},
    {"pattern": "改天册", "replacement": "改用天册", "note": "替加环素品牌名前缺动词"},
    {"pattern": "性别冬", "replacement": "性别女", "note": "性别女误识别"}
  ]
}

# PDF Processor - 依赖（PDF 处理功能）

# PDF 转图像
pdf2image>=1.17.0

# 图像处理
opencv-python>=4.9.0

# 图像操作
pillow>=10.0.0

# 数值计算
numpy>=1.24.0

# PDF 文字层叠加、解密、水印检测、压缩
pymupdf>=1.23.0

# PDF 解密、合并辅助
pypdf>=3.17.0

# 本地 Paddle 双层实验后端（非默认生产链路，需要时再安装）
# paddleocr
# paddlepaddle

"""
OCR dump/resume 模块 + Agent 修正。

工作流：
  Step 1: pdf-ocr.py -i input.pdf -o output.pdf --ocr-dump /tmp/ocr_dump.json
          → OCR 识别，结果存入 JSON，同时生成可读文本，不生成 PDF
  Step 2: agent 审查可读文本，发现 OCR 错误
          → 直接编辑 dump JSON 文件中的 text 字段，或生成 corrections 文件
  Step 3: pdf-ocr.py -i input.pdf -o output.pdf --ocr-resume /tmp/ocr_dump.json
          → 加载 dump（含 agent 修正）+ 生成双层 PDF
"""

from __future__ import annotations

import json
import re
from datetime import datetime
from pathlib import Path

_SKILL_ROOT = Path(__file__).resolve().parent.parent
ARCHIVE_DIR = _SKILL_ROOT / "archive"


# ---------- 序列化 / 反序列化 ----------

def _entries_to_serializable(entries: list[dict]) -> list[dict]:
    """将 page_entries 转为 JSON 可序列化格式。"""
    result = []
    for entry in entries:
        rows = []
        for text, score, poly in entry.get("rows", []):
            rows.append({
                "text": text,
                "score": score,
                "poly": poly,
            })
        result.append({
            "rows": rows,
            "width": entry.get("width"),
            "height": entry.get("height"),
        })
    return result


def _entries_from_serializable(data: list[dict]) -> list[dict]:
    """从 JSON 格式恢复 page_entries（含 tuple rows）。"""
    entries = []
    for page in data:
        rows = []
        for block in page.get("rows", []):
            text = block.get("text", "")
            score = block.get("score", 1.0)
            poly = block.get("poly", [[0, 0], [0, 0], [0, 0], [0, 0]])
            rows.append((text, score, poly))
        entries.append({
            "rows": rows,
            "width": page.get("width"),
            "height": page.get("height"),
        })
    return entries


def dump_page_entries(
    entries: list[dict],
    path: str | Path,
    source: str = "",
    model: str = "",
) -> None:
    """将 OCR 结果保存为 JSON 文件，供 agent 审查。"""
    p = Path(path)
    data = {
        "source": source,
        "model": model,
        "total_pages": len(entries),
        "pages": _entries_to_serializable(entries),
    }

    p.parent.mkdir(parents=True, exist_ok=True)
    with open(p, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def load_page_entries(path: str | Path) -> tuple[list[dict], dict]:
    """从 dump 文件加载 OCR 结果。

    Returns:
        (page_entries, metadata_dict)
    """
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(f"OCR dump 文件不存在: {p}")

    with open(p, "r", encoding="utf-8") as f:
        data = json.load(f)

    pages = data.get("pages", [])
    entries = _entries_from_serializable(pages)
    metadata = {
        "source": data.get("source", ""),
        "model": data.get("model", ""),
    }
    return entries, metadata


# ---------- 可读文本生成 ----------

def generate_readable_text(entries: list[dict]) -> str:
    """从 page_entries 生成连续可读文本，同段落行自动拼接，段落间空行分隔。"""
    paragraphs: list[str] = []
    current: list[str] = []
    had_wrap = False
    prev_y_bot: float | None = None
    prev_x_left: float = 0.0
    prev_h: float = 0.0

    for entry in entries:
        for text, score, poly in entry.get("rows", []):
            content = text.strip()
            if not content or score <= 0:
                continue

            if not poly or len(poly) < 4:
                if current:
                    paragraphs.append("".join(current))
                    current = []
                paragraphs.append(content)
                prev_y_bot = None
                had_wrap = False
                continue

            ys = [p[1] for p in poly]
            xs = [p[0] for p in poly]
            y_top = min(ys)
            y_bot = max(ys)
            x_left = min(xs)
            h = max(y_bot - y_top, 1.0)

            if prev_y_bot is None:
                current = [content]
                had_wrap = False
                prev_y_bot = y_bot
                prev_x_left = x_left
                prev_h = h
                continue

            y_gap = y_top - prev_y_bot
            min_h = max(min(prev_h, h), 1.0)
            x_diff = x_left - prev_x_left

            # 段落边界：负间距（同行重叠）、大间距（跨段落）
            if y_gap <= 0 or y_gap >= min_h:
                if current:
                    paragraphs.append("".join(current))
                current = [content]
                had_wrap = False
            else:
                is_wrap = x_diff < -min_h * 0.3
                is_same_margin = abs(x_diff) < min_h * 0.15

                if is_wrap:
                    current.append(content)
                    had_wrap = True
                elif is_same_margin and had_wrap:
                    current.append(content)
                else:
                    if current:
                        paragraphs.append("".join(current))
                    current = [content]
                    had_wrap = False

            prev_y_bot = y_bot
            prev_x_left = x_left
            prev_h = h

    if current:
        paragraphs.append("".join(current))

    return "\n\n".join(paragraphs)


# ---------- Agent 修正（from/to 文本替换）----------

def load_agent_corrections(path: str | Path) -> list[tuple[str, str]]:
    """加载 agent 修正文件。

    JSON 格式：
    [
      {"from": "性别芝", "to": "性别女"},
      {"from": "沙库巴去缬沙坦", "to": "沙库巴曲缬沙坦"}
    ]
    """
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(f"修正文件不存在: {p}")

    with open(p, "r", encoding="utf-8") as f:
        data = json.load(f)

    if not isinstance(data, list):
        raise ValueError("修正文件格式错误：期望 JSON 数组 [{from, to}, ...]")

    corrections = []
    for item in data:
        if isinstance(item, dict):
            src = item.get("from")
            dst = item.get("to")
            if src and dst:
                corrections.append((str(src), str(dst)))

    return corrections


def apply_agent_corrections(
    entries: list[dict],
    corrections: list[tuple[str, str]],
    quiet: bool = False,
) -> tuple[list[dict], int]:
    """对 page entries 应用 agent 的 from/to 文本替换。"""
    total = 0
    for entry in entries:
        if "rows" not in entry:
            continue
        new_rows = []
        for text, score, poly in entry["rows"]:
            corrected = text
            for src, dst in corrections:
                if src in corrected:
                    corrected = corrected.replace(src, dst)
            if corrected != text:
                total += 1
            new_rows.append((corrected, score, poly))
        entry["rows"] = new_rows

    if total and not quiet:
        print(f"  Agent 修正: {total} 处文本已纠正")

    return entries, total


# ---------- Archive 归档（仅保存 MD）----------

def _sanitize_name(name: str) -> str:
    """清理文件名中的特殊字符，保留中文、字母、数字、下划线和连字符。"""
    name = re.sub(r'[\\/:*?"<>|]', '_', name)
    name = re.sub(r'[“”‘’]', '', name)  # 中文引号
    name = re.sub(r'\s+', ' ', name).strip()
    return name or "unnamed"


def archive_ocr_result(
    source_path: str | Path,
    entries: list[dict],
    model: str = "",
    backend: str = "",
    extra_meta: dict | None = None,
    original_source_path: str | Path | None = None,
) -> tuple[Path, Path]:
    """将 OCR 可读文本归档到 skill 内部 archive/，同时输出到原文件同目录。

    归档格式：archive/YYYYMMDD_HHMMSS_文件名/原文件名.md + conversion_meta.json
    同时在原文件同目录生成同名 .md 文件。

    Args:
        source_path: 实际处理的 PDF 路径（可能是临时文件）。
        original_source_path: 用户最初传入的原始 PDF 路径。
            如果提供，归档目录名和伴生 .md 以原始文件为准，
            source_file 记录原始路径，working_file 记录实际处理路径。

    Returns:
        (archive_md_path, source_dir_md_path)
    """
    working = Path(source_path)
    original = Path(original_source_path) if original_source_path else None

    # 归档命名和伴生 .md 以原始文件为准
    naming_source = original or working
    stem = _sanitize_name(naming_source.stem)
    md_name = naming_source.stem + ".md"
    md_content = generate_readable_text(entries)

    # archive 内部归档
    now = datetime.now()
    timestamp = now.strftime("%Y%m%d_%H%M%S")
    archive_path = ARCHIVE_DIR / f"{timestamp}_{stem}"
    archive_path.mkdir(parents=True, exist_ok=True)
    archive_md = archive_path / md_name
    archive_md.write_text(md_content, encoding="utf-8")

    # 运行记录 JSON
    total_pages = len(entries)
    total_rows = sum(len(e.get("rows", [])) for e in entries)
    meta = {
        "timestamp": now.isoformat(),
        "source_file": str(naming_source),
        "archive_path": str(archive_path),
        "model": model,
        "backend": backend,
        "total_pages": total_pages,
        "total_text_blocks": total_rows,
    }
    if original and original != working:
        meta["working_file"] = str(working)
    if extra_meta:
        meta.update(extra_meta)
    meta_path = archive_path / "conversion_meta.json"
    with open(meta_path, "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False, indent=2)

    # 原文件同目录输出（始终写到原始文件旁边）
    source_md = naming_source.parent / md_name
    source_md.write_text(md_content, encoding="utf-8")

    return archive_md, source_md


def archive_preprocess_result(
    source_path: str | Path,
    preprocess_meta: dict | None = None,
    output_path: str | Path | None = None,
    original_source_path: str | Path | None = None,
) -> Path:
    """仅预处理模式的轻量归档（无 OCR 文本，只记录元数据）。

    归档格式：archive/YYYYMMDD_HHMMSS_文件名/conversion_meta.json

    Args:
        source_path: 实际处理的 PDF 路径。
        preprocess_meta: 预处理参数和统计信息。
        output_path: 最终输出文件路径。
        original_source_path: 用户最初传入的原始 PDF 路径。

    Returns:
        archive_dir_path
    """
    working = Path(source_path)
    original = Path(original_source_path) if original_source_path else None
    naming_source = original or working
    stem = _sanitize_name(naming_source.stem)

    now = datetime.now()
    timestamp = now.strftime("%Y%m%d_%H%M%S")
    archive_path = ARCHIVE_DIR / f"{timestamp}_{stem}"
    archive_path.mkdir(parents=True, exist_ok=True)

    meta = {
        "timestamp": now.isoformat(),
        "source_file": str(naming_source),
        "archive_path": str(archive_path),
        "mode": "preprocess_only",
    }
    if output_path:
        meta["output_file"] = str(Path(output_path))
    if original and original != working:
        meta["working_file"] = str(working)
    if preprocess_meta:
        meta["preprocess"] = preprocess_meta

    meta_path = archive_path / "conversion_meta.json"
    with open(meta_path, "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False, indent=2)

    return archive_path

#!/usr/bin/env python3
"""
PDF 双层叠层核心模块。

将 OCR 结果（文字 + 坐标）叠入 PDF 透明文字层。
支持两种来源：
1. 本地 PaddleOCR predict 输出
2. 外部 API 返回的 payload

公共函数：
- normalize_cjk_spacing
- page_has_text_layer
- parse_paddle_predict_result
- calculate_font_size
- extract_page_image_size
- extract_page_entries_from_api_payload
- infer_page_scale
- apply_page_entries_as_layered_pdf
- apply_api_payload_as_layered_pdf
- save_output_from_api_payload
"""

from __future__ import annotations

import base64
import os
import platform
import re
import shutil
import subprocess
from datetime import datetime
from pathlib import Path

from pdf_runtime import http_get_bytes


# ---------- 通用 Payload 提取 ----------

def extract_payload(resp: dict) -> dict:
    """
    提取有效载荷。
    支持两类格式：
    1) 顶层直接放 output 字段
    2) `data` 字段中放 output 字段
    """
    result = resp.get("result")
    if isinstance(result, dict):
        return result
    if isinstance(result, list):
        return {"ocrResults": result}

    data = resp.get("data")
    if isinstance(data, dict):
        return data
    if isinstance(data, list):
        return {"ocrResults": data}
    return resp


# ---------- CJK 空格归一化 ----------

_CJK_SPACE_RE = re.compile(
    r"(?<=[㐀-䶿一-鿿豈-﫿　-〿＀-￯])"
    r"\s+"
    r"(?=[㐀-䶿一-鿿豈-﫿　-〿＀-￯])"
)
_CJK_BEFORE_PUNC_SPACE_RE = re.compile(r"\s+([，。！？；：、）》】」』）])")
_CJK_AFTER_OPEN_PUNC_SPACE_RE = re.compile(r"([（《【「『])\s+")


def normalize_cjk_spacing(text: str) -> str:
    """移除 CJK 字符间误插入空格，保留英文词间空格。"""
    if not text:
        return text
    text = _CJK_SPACE_RE.sub("", text)
    text = _CJK_BEFORE_PUNC_SPACE_RE.sub(r"\1", text)
    text = _CJK_AFTER_OPEN_PUNC_SPACE_RE.sub(r"\1", text)
    text = re.sub(r"\s{2,}", " ", text)
    return text.strip()


# ---------- 几何工具 ----------

def _poly_to_points(poly) -> list[list[float]]:
    """将多种 polygon 表示统一为 [[x,y], ...]。"""
    if poly is None:
        return []
    if not isinstance(poly, (list, tuple)):
        return []
    if not poly:
        return []

    first = poly[0]
    if isinstance(first, (list, tuple)):
        out = []
        for p in poly:
            if not isinstance(p, (list, tuple)) or len(p) < 2:
                continue
            try:
                out.append([float(p[0]), float(p[1])])
            except Exception:
                continue
        return out

    # 扁平数组: [x1,y1,x2,y2,...]
    out = []
    if len(poly) >= 8:
        for i in range(0, len(poly) - 1, 2):
            try:
                out.append([float(poly[i]), float(poly[i + 1])])
            except Exception:
                continue
    return out


def _as_float(val, default: float = 0.0) -> float:
    try:
        return float(val)
    except Exception:
        return default


def _as_num_list(obj, length: int) -> list[float] | None:
    if not isinstance(obj, (list, tuple)) or len(obj) < length:
        return None
    nums = []
    for i in range(length):
        try:
            nums.append(float(obj[i]))
        except Exception:
            return None
    return nums


def _bbox_to_poly4(bbox) -> list[list[float]]:
    nums = _as_num_list(bbox, 4)
    if not nums:
        return []
    x0, y0, x1, y1 = nums
    return [[x0, y0], [x1, y0], [x1, y1], [x0, y1]]


# ---------- 页面判断 ----------

def page_has_text_layer(page, min_chars: int) -> bool:
    """判断页面是否已有较明显文本层。"""
    text = page.get_text("text")
    return len(text.strip()) >= min_chars


# ---------- PaddleOCR 结果解析 ----------

def _parse_rec_dict(item: dict) -> list[tuple[str, float, list[list[float]]]]:
    texts = item.get("rec_texts")
    scores = item.get("rec_scores")
    polys = item.get("rec_polys")

    if texts is None and "texts" in item:
        texts = item.get("texts")
    if scores is None and "scores" in item:
        scores = item.get("scores")
    if polys is None:
        polys = (
            item.get("dt_polys")
            or item.get("rec_boxes")
            or item.get("polys")
            or item.get("text_region")
        )

    if not isinstance(texts, list) or not isinstance(polys, list):
        return []

    if not isinstance(scores, list):
        scores = []

    parsed = []
    for i, text in enumerate(texts):
        poly = polys[i] if i < len(polys) else None
        poly4 = _poly_to_points(poly)
        if len(poly4) < 4:
            continue
        score = _as_float(scores[i], default=1.0) if i < len(scores) else 1.0
        parsed.append((str(text), score, poly4[:4]))
    return parsed


def parse_paddle_predict_result(result) -> list[tuple[str, float, list[list[float]]]]:
    """
    解析 PaddleOCR 输出，统一为: [(text, score, poly4), ...]
    poly4: [[x1,y1],[x2,y2],[x3,y3],[x4,y4]]
    """
    if result is None:
        return []

    # 允许直接传 dict（常见于 API 的 prunedResult）
    if isinstance(result, dict):
        parsed = _parse_rec_dict(result)
        if parsed:
            return parsed

        for key in ("prunedResult", "ocrResult", "result", "data"):
            nested = result.get(key)
            nested_rows = parse_paddle_predict_result(nested)
            if nested_rows:
                return nested_rows
        return []

    # 新版 `predict` 常见结构: [ {rec_texts, rec_scores, rec_polys, ...} ]
    if isinstance(result, list) and result and isinstance(result[0], dict):
        rows = []
        for item in result:
            rows.extend(parse_paddle_predict_result(item))
        return rows

    # 旧版结构: [ [poly, (text, score)], ... ]
    if isinstance(result, list) and result and isinstance(result[0], list):
        blocks = result[0]
        parsed = []
        if isinstance(blocks, list):
            for block in blocks:
                if not isinstance(block, list) or len(block) < 2:
                    continue
                poly = block[0]
                rec = block[1]
                if not isinstance(rec, (list, tuple)) or len(rec) < 2:
                    continue
                text = str(rec[0])
                score = _as_float(rec[1], default=1.0)
                poly4 = _poly_to_points(poly)
                if len(poly4) >= 4:
                    parsed.append((text, score, poly4[:4]))
        return parsed

    return []


# ---------- 字号计算 ----------

def calculate_font_size(font, text: str, w: float, h: float) -> float:
    """计算贴合文本框的字号，支持多行文本。

    策略：
    1. 估算文本在当前框内需要的行数
    2. 按行高算出初始字号
    3. 微调使每行尽量填满宽度
    4. 字号不超过框高
    """
    if not text:
        return max(1.0, h)
    if h > w:
        w, h = h, w
    min_size = 5.0
    if w <= 0 or h <= 0:
        return min_size

    def text_len(size: float) -> float:
        return font.text_length(text, fontsize=size)

    # 估算行数：在字号=h 时文本需要几行
    text_w_at_h = text_len(h)
    if text_w_at_h <= 0:
        return max(min_size, h)
    estimated_lines = max(1.0, text_w_at_h / w)

    # 按行数算出单行高度
    line_height = h / estimated_lines
    fontsize = max(min_size, round(line_height))

    # 微调：如果单行文字比框宽，缩小字号
    if text_len(fontsize) > w * estimated_lines * 1.1:
        while text_len(fontsize) > w * estimated_lines * 1.1 and fontsize > min_size:
            fontsize -= 0.5

    # 单行文本：字号不超过框高
    if estimated_lines <= 1.0:
        max_size = max(min_size, h)
        if fontsize > max_size:
            fontsize = max_size

    return max(min_size, fontsize)


# ---------- 页面图像尺寸 ----------

def _pick_positive_number(value):
    try:
        num = float(value)
    except Exception:
        return None
    return num if num > 0 else None


def extract_page_image_size(*objs) -> tuple[float | None, float | None]:
    """从页面结果中提取 OCR 坐标空间的宽高。"""
    key_pairs = [
        ("imageWidth", "imageHeight"),
        ("image_width", "image_height"),
        ("imgW", "imgH"),
        ("img_w", "img_h"),
        ("pageWidth", "pageHeight"),
        ("page_width", "page_height"),
        ("width", "height"),
    ]
    shape_keys = (
        "imageShape",
        "image_shape",
        "inputImageShape",
        "input_image_shape",
        "shape",
    )

    for obj in objs:
        if not isinstance(obj, dict):
            continue
        for wk, hk in key_pairs:
            if wk in obj and hk in obj:
                w = _pick_positive_number(obj.get(wk))
                h = _pick_positive_number(obj.get(hk))
                if w and h:
                    return w, h

        for sk in shape_keys:
            shape = obj.get(sk)
            if isinstance(shape, (list, tuple)) and len(shape) >= 2:
                h = _pick_positive_number(shape[0])
                w = _pick_positive_number(shape[1])
                if w and h:
                    return w, h
            if isinstance(shape, dict):
                w = _pick_positive_number(shape.get("w") or shape.get("width"))
                h = _pick_positive_number(shape.get("h") or shape.get("height"))
                if w and h:
                    return w, h

    return None, None


# ---------- API payload -> 分页 entries ----------

def extract_page_entries_from_api_payload(payload: dict) -> list[dict]:
    """
    从外部 API 返回中提取分页 OCR 结果。
    返回: [{"rows": [...], "width": x, "height": y}, ...]
    """
    page_sources = None
    if isinstance(payload, list):
        page_sources = payload
    elif isinstance(payload, dict):
        for key in ("ocrResults", "layoutParsingResults", "pageResults", "pages", "results"):
            if isinstance(payload.get(key), list):
                page_sources = payload.get(key)
                break
        if page_sources is None and (
            isinstance(payload.get("prunedResult"), dict) or "rec_texts" in payload
        ):
            page_sources = [payload]

    if not isinstance(page_sources, list):
        return []

    entries = []
    for page_obj in page_sources:
        pruned = None
        if isinstance(page_obj, dict):
            pruned = page_obj.get("prunedResult")

        candidate = pruned if isinstance(pruned, dict) else page_obj
        rows = parse_paddle_predict_result(candidate)
        width, height = extract_page_image_size(page_obj, pruned)
        entries.append(
            {
                "rows": rows,
                "width": width,
                "height": height,
            }
        )
    return entries


# ---------- 缩放推断 ----------

def infer_page_scale(page_rect, rows, source_w, source_h) -> tuple[float, float]:
    """推断 OCR 坐标到 PDF 页面坐标的缩放系数。"""
    if source_w and source_h:
        return page_rect.width / source_w, page_rect.height / source_h

    max_x = 0.0
    max_y = 0.0
    for _, _, poly in rows:
        for p in poly:
            try:
                max_x = max(max_x, float(p[0]))
                max_y = max(max_y, float(p[1]))
            except Exception:
                continue

    if max_x <= page_rect.width * 1.25 and max_y <= page_rect.height * 1.25:
        return 1.0, 1.0

    if max_x <= 0 or max_y <= 0:
        return 1.0, 1.0
    return page_rect.width / max_x, page_rect.height / max_y


# ---------- 透明文字块插入（去重后的共享函数） ----------

def _insert_text_blocks(
    page,
    font,
    rows: list[tuple[str, float, list[list[float]]]],
    *,
    scale_x: float,
    scale_y: float,
    min_score: float,
    cjk_normalize: bool,
    page_rotation: int,
    source_name: str,
    pno: int,
    total_pages: int,
    quiet: bool,
) -> int:
    import fitz
    """
    向单个 PDF 页面插入透明文字块。

    Returns:
        插入的文本块数量。
    """
    font_inserted = False
    page_inserted = 0

    page.clean_contents()

    for text, score, poly in rows:
        if score < min_score:
            continue
        content = text.strip()
        if not content:
            continue
        if cjk_normalize:
            content = normalize_cjk_spacing(content)
            if not content:
                continue

        xs = [p[0] * scale_x for p in poly]
        ys = [p[1] * scale_y for p in poly]

        x0 = min(xs)
        x1 = max(xs)
        y0 = min(ys)
        y1 = max(ys)

        w = max(1.0, x1 - x0)
        h = max(1.0, y1 - y0)
        fontsize = calculate_font_size(font, content, w, h)

        if not font_inserted:
            page.insert_font(fontname="cjk", fontbuffer=font.buffer)
            font_inserted = True

        point = fitz.Point(x0, y1) * page.derotation_matrix
        page.insert_text(
            point,
            content,
            fontsize=fontsize,
            fontname="cjk",
            rotate=page_rotation,
            stroke_opacity=0,
            fill_opacity=0,
            render_mode=3,
        )
        page_inserted += 1

    if not quiet:
        print(f"  第 {pno}/{total_pages} 页({source_name}): 新增 {page_inserted} 文本块")

    return page_inserted


# ---------- 叠层 PDF（分页 entries） ----------

def apply_page_entries_as_layered_pdf(page_entries: list[dict], args, source_name: str) -> bool:
    """将分页 OCR 结果叠层为双层 PDF。"""
    if not page_entries:
        return False

    import fitz

    doc = fitz.open(args.input)
    font = fitz.Font("cjk")

    inserted_pages = 0
    inserted_blocks = 0
    skipped_pages = 0
    total_pages = len(doc)

    cjk_normalize = not args.no_paddle_cjk_space_normalize

    for pno, page in enumerate(doc, start=1):
        if pno - 1 >= len(page_entries):
            continue

        if args.mode == "skip" and page_has_text_layer(page, args.paddle_skip_text_min_chars):
            skipped_pages += 1
            continue

        entry = page_entries[pno - 1]
        rows = entry.get("rows") or []
        if not rows:
            continue

        scale_x, scale_y = infer_page_scale(
            page.rect,
            rows,
            entry.get("width"),
            entry.get("height"),
        )

        page_rotation = int(page.rotation) if page.rotation else 0

        page_inserted = _insert_text_blocks(
            page,
            font,
            rows,
            scale_x=scale_x,
            scale_y=scale_y,
            min_score=args.paddle_min_score,
            cjk_normalize=cjk_normalize,
            page_rotation=page_rotation,
            source_name=source_name,
            pno=pno,
            total_pages=total_pages,
            quiet=args.quiet,
        )

        if page_inserted > 0:
            inserted_pages += 1
            inserted_blocks += page_inserted

    if inserted_pages == 0:
        doc.close()
        src = Path(args.input).resolve()
        dst = Path(args.output).resolve()
        if src != dst:
            shutil.copy2(src, dst)
        if not args.quiet:
            print(f"{source_name} 已返回 OCR 结果，但无可叠层文本块，已原样输出。")
        return True

    try:
        doc.subset_fonts()
    except Exception:
        pass

    doc.save(
        args.output,
        garbage=4,
        clean=1, deflate=1, deflate_images=1, deflate_fonts=1,
        use_objstms=1, compression_effort=100,
    )
    doc.close()

    # 保留原文件时间戳（创建时间 + 修改时间）
    try:
        src_stat = Path(args.input).stat()
        mtime = src_stat.st_mtime
        birthtime = getattr(src_stat, "st_birthtime", mtime)
        os.utime(args.output, (mtime, mtime))
        if platform.system() == "Darwin":
            try:
                dt = datetime.fromtimestamp(birthtime)
                date_str = dt.strftime("%m/%d/%Y %H:%M:%S")
                subprocess.run(
                    ["SetFile", "-d", date_str, str(args.output)],
                    check=True, capture_output=True,
                )
            except (subprocess.CalledProcessError, FileNotFoundError):
                pass
    except Exception:
        pass

    if not args.quiet:
        print(f"\n{source_name} OCR 叠层完成:")
        print(f"  新增页面: {inserted_pages}/{total_pages}")
        print(f"  新增文本块: {inserted_blocks}")
        print(f"  跳过页面: {skipped_pages}")
    return True


# ---------- 叠层 PDF（API payload） ----------

def apply_api_payload_as_layered_pdf(payload: dict, args) -> bool:
    """将 API 返回的 OCR 结果叠层为双层 PDF。"""
    page_entries = extract_page_entries_from_api_payload(payload)
    return apply_page_entries_as_layered_pdf(page_entries, args, source_name="API")


# ---------- 保存 API 输出 PDF ----------

def save_output_from_api_payload(payload: dict, output_path: Path, timeout: int):
    """从 API payload 保存输出 PDF。"""
    if "output_pdf_base64" in payload:
        raw = base64.b64decode(payload["output_pdf_base64"])
        output_path.write_bytes(raw)
        return

    if "output_pdf_url" in payload:
        raw = http_get_bytes(payload["output_pdf_url"], timeout=timeout)
        output_path.write_bytes(raw)
        return

    if "output_pdf_path" in payload:
        p = Path(payload["output_pdf_path"])
        if not p.exists():
            raise RuntimeError(f"API 返回的输出路径不存在: {p}")
        shutil.copy2(p, output_path)
        return

    raise RuntimeError(
        "API 响应未直接返回 PDF。可继续尝试解析 OCR 结果并本地叠层。"
    )

#!/usr/bin/env python3
"""Historical local PaddleOCR layered-PDF backend.

This module is intentionally kept outside the public OCR entrypoint. The
production path prefers external PaddleOCR / MinerU APIs and local ocrmypdf
fallback, but this local implementation is preserved for future hardware-rich
environments and controlled experiments.
"""

from __future__ import annotations

import os
import platform
import shutil
from pathlib import Path

from pdf_runtime import print_dependency_help
from pdf_ocr_layered import (
    _insert_text_blocks,
    page_has_text_layer,
    parse_paddle_predict_result,
)

try:
    import fitz  # PyMuPDF
    HAS_PYMUPDF = True
except Exception:
    HAS_PYMUPDF = False

try:
    import numpy as np
    HAS_NUMPY = True
except Exception:
    HAS_NUMPY = False

PaddleOCR = None
HAS_PADDLEOCR = None


PADDLE_PROFILE_PRESETS = {
    "quality": {
        "det_model": "PP-OCRv5_server_det",
        "rec_model": "PP-OCRv5_server_rec",
        "dpi": 300,
        "det_limit_side_len": 1536,
    },
    "balanced": {
        "det_model": "PP-OCRv5_mobile_det",
        "rec_model": "PP-OCRv5_mobile_rec",
        "dpi": 300,
        "det_limit_side_len": 1216,
    },
    "speed": {
        "det_model": "PP-OCRv5_mobile_det",
        "rec_model": "PP-OCRv5_mobile_rec",
        "dpi": 260,
        "det_limit_side_len": 960,
    },
}


def get_paddle_layered_missing_dependencies() -> list[str]:
    """Detect Python dependencies required by the local Paddle backend."""
    global PaddleOCR, HAS_PADDLEOCR

    missing = []
    if not HAS_PYMUPDF:
        missing.append("pymupdf")
    if not HAS_NUMPY:
        missing.append("numpy")
    if HAS_PADDLEOCR is None:
        try:
            from paddleocr import PaddleOCR as _PaddleOCR  # lazy import
            PaddleOCR = _PaddleOCR
            HAS_PADDLEOCR = True
        except Exception:
            HAS_PADDLEOCR = False
    if not HAS_PADDLEOCR:
        missing.extend(["paddleocr", "paddlepaddle"])
    return missing


def ensure_paddle_layered_available() -> bool:
    """Check local Paddle backend dependencies."""
    missing = get_paddle_layered_missing_dependencies()
    if missing:
        print_dependency_help(
            "本地 Paddle 双层引擎",
            missing_python=missing,
            install_commands=[
                "pip install pymupdf numpy paddleocr paddlepaddle",
            ],
            extra_notes=[
                "本地 Paddle 后端是历史实验能力，不属于默认生产链路。",
                "首次运行 Paddle 相关能力时，模型初始化也可能继续耗时。",
            ],
        )
        return False
    return True


def map_tesseract_lang_to_paddle(language: str) -> str:
    """Map Tesseract-style language values to PaddleOCR language values."""
    lang = (language or "").lower()
    if "chi_sim" in lang or "chi_tra" in lang or "ch" in lang:
        return "ch"
    if "jpn" in lang or "jap" in lang:
        return "japan"
    if "kor" in lang:
        return "korean"
    if "eng" in lang:
        return "en"
    return "ch"


def detect_cuda_device_count() -> int:
    """Detect available CUDA devices on Windows/Linux."""
    system = platform.system().lower()
    if system not in {"windows", "linux"}:
        return 0

    try:
        import paddle  # type: ignore
    except Exception:
        return 0

    try:
        device_mod = getattr(paddle, "device", None)
        if device_mod is None:
            return 0

        is_cuda_fn = getattr(device_mod, "is_compiled_with_cuda", None)
        if not callable(is_cuda_fn) or not is_cuda_fn():
            return 0

        cuda_mod = getattr(device_mod, "cuda", None)
        if cuda_mod is None:
            return 1

        count_fn = getattr(cuda_mod, "device_count", None)
        if callable(count_fn):
            count = int(count_fn() or 0)
            return max(0, count)

        return 1
    except Exception:
        return 0


def resolve_paddle_device(args) -> tuple[bool, str]:
    """Resolve local Paddle runtime device."""
    if args.paddle_use_gpu:
        return True, "用户显式指定 GPU"

    cuda_count = detect_cuda_device_count()
    if cuda_count > 0:
        return True, f"{platform.system()} 自动检测到 CUDA GPU x{cuda_count}"

    return False, f"{platform.system()} 默认 CPU"


def resolve_paddle_profile(args, total_pages: int, use_gpu: bool) -> tuple[str, str]:
    """Resolve local Paddle profile."""
    requested = args.paddle_profile
    if requested != "auto":
        return requested, "用户显式指定"

    if args.paddle_long_doc_pages > 0 and total_pages >= args.paddle_long_doc_pages:
        return "speed", f"页数 {total_pages} >= 阈值 {args.paddle_long_doc_pages}"

    if use_gpu:
        return "quality", "启用 GPU"

    system = platform.system().lower()
    if system in {"darwin", "windows"}:
        return "balanced", f"{platform.system()} CPU 默认平衡档"

    return "quality", f"{platform.system()} 默认高精度档"


def build_paddle_runtime_config(args, total_pages: int, use_gpu: bool) -> dict:
    """Build final local Paddle runtime config from profile and CLI args."""
    profile, reason = resolve_paddle_profile(args, total_pages, use_gpu)
    preset = PADDLE_PROFILE_PRESETS[profile]

    det_model = args.paddle_det_model_name or preset["det_model"]
    rec_model = args.paddle_rec_model_name or preset["rec_model"]
    dpi = args.paddle_dpi if args.paddle_dpi_user_set else preset["dpi"]
    det_limit = (
        args.paddle_det_limit_side_len
        if args.paddle_det_limit_user_set
        else preset["det_limit_side_len"]
    )

    return {
        "profile": profile,
        "profile_reason": reason,
        "use_gpu": bool(use_gpu),
        "det_model": det_model,
        "rec_model": rec_model,
        "dpi": int(dpi),
        "det_limit_side_len": int(det_limit),
    }


def run_local_paddle_layered_backend(args, fallback_backend=None):
    """Run local PaddleOCR + PyMuPDF transparent text-layer generation."""
    if not args.keep_paddle_model_source_check:
        os.environ.setdefault("PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK", "True")
        os.environ.setdefault("DISABLE_MODEL_SOURCE_CHECK", "True")

    if args.paddle_model_source:
        os.environ["PADDLE_PDX_MODEL_SOURCE"] = args.paddle_model_source.lower()

    if not ensure_paddle_layered_available():
        raise RuntimeError("本地 Paddle 双层引擎不可用")

    paddle_lang = args.paddle_lang or map_tesseract_lang_to_paddle(args.language)

    with fitz.open(args.input) as probe_doc:
        total_pages = len(probe_doc)
        has_text_pages = [
            i + 1
            for i, page in enumerate(probe_doc)
            if page_has_text_layer(page, args.paddle_skip_text_min_chars)
        ]

    use_gpu, device_reason = resolve_paddle_device(args)
    runtime_cfg = build_paddle_runtime_config(args, total_pages, use_gpu)

    if not args.quiet:
        print("\n本地 Paddle 双层后端参数:")
        print(f"  profile: {runtime_cfg['profile']} ({runtime_cfg['profile_reason']})")
        print(f"  device: {'gpu' if runtime_cfg['use_gpu'] else 'cpu'} ({device_reason})")
        print(f"  lang: {paddle_lang}")
        print(f"  dpi: {runtime_cfg['dpi']}")
        print(f"  min_score: {args.paddle_min_score}")
        print(f"  skip_text_min_chars: {args.paddle_skip_text_min_chars}")
        print(f"  textline_orientation: {args.paddle_textline_orientation}")
        print(f"  det_limit_side_len: {runtime_cfg['det_limit_side_len']}")
        print(f"  det_model: {runtime_cfg['det_model']}")
        print(f"  rec_model: {runtime_cfg['rec_model']}")
        print(f"  model_source_check: {'keep' if args.keep_paddle_model_source_check else 'disabled'}")
        if args.paddle_model_source:
            print(f"  model_source: {args.paddle_model_source.lower()}")
        print(f"  normalize_cjk_spaces: {not args.no_paddle_cjk_space_normalize}")

    if args.dry_run:
        print("[DRY-RUN] 本地 Paddle 双层后端参数已输出，未实际执行。")
        args.backend_used = "local_paddle_layered"
        return

    if has_text_pages and args.mode in {"redo", "force"}:
        if fallback_backend is None:
            raise RuntimeError("检测到已有文本层，local_paddle_layered 需要 ocrmypdf 兜底处理 redo/force")
        if not args.quiet:
            print(
                "警告: 检测到已有文本层，`redo/force` 语义下自动回退 ocrmypdf，"
                "以避免重复文字层或旧层残留。"
            )
        fallback_backend(args)
        return

    ocr_kwargs = {
        "lang": paddle_lang,
        "use_doc_orientation_classify": False,
        "use_doc_unwarping": False,
        "use_textline_orientation": args.paddle_textline_orientation,
        "text_det_limit_side_len": runtime_cfg["det_limit_side_len"],
        "text_det_limit_type": "max",
        "text_detection_model_name": runtime_cfg["det_model"],
        "text_recognition_model_name": runtime_cfg["rec_model"],
    }
    if runtime_cfg["use_gpu"]:
        ocr_kwargs["device"] = "gpu"

    try:
        ocr = PaddleOCR(**ocr_kwargs)
    except TypeError:
        ocr_kwargs.pop("device", None)
        ocr = PaddleOCR(**ocr_kwargs)
    except Exception as e:
        if runtime_cfg["use_gpu"] and not args.quiet:
            print(f"警告: PaddleOCR GPU 初始化失败，回退 CPU。原因: {e}")
            ocr_kwargs.pop("device", None)
            ocr = PaddleOCR(**ocr_kwargs)
        else:
            raise

    doc = fitz.open(args.input)
    font = fitz.Font("cjk")
    inserted_pages = 0
    inserted_blocks = 0
    skipped_pages = 0
    cjk_normalize = not args.no_paddle_cjk_space_normalize

    for pno, page in enumerate(doc, start=1):
        if args.mode == "skip" and page_has_text_layer(page, args.paddle_skip_text_min_chars):
            skipped_pages += 1
            continue

        pix = page.get_pixmap(dpi=runtime_cfg["dpi"], alpha=False)
        arr = np.frombuffer(pix.samples, dtype=np.uint8)
        arr = arr.reshape(pix.height, pix.width, pix.n)
        if pix.n == 4:
            arr = arr[:, :, :3]

        rows = parse_paddle_predict_result(ocr.predict(arr))
        if not rows:
            continue

        page_inserted = _insert_text_blocks(
            page,
            font,
            rows,
            scale_x=page.rect.width / float(pix.width),
            scale_y=page.rect.height / float(pix.height),
            min_score=args.paddle_min_score,
            cjk_normalize=cjk_normalize,
            page_rotation=int(page.rotation) if page.rotation else 0,
            source_name="Paddle",
            pno=pno,
            total_pages=total_pages,
            quiet=args.quiet,
        )

        if page_inserted > 0:
            inserted_pages += 1
            inserted_blocks += page_inserted

    if inserted_pages == 0:
        doc.close()
        src = Path(args.input).resolve()
        dst = Path(args.output).resolve()
        if src != dst:
            shutil.copy2(src, dst)
        if not args.quiet:
            print("未新增 OCR 文本层，已原样输出。")
        args.backend_used = "local_paddle_layered"
        return

    try:
        doc.subset_fonts()
    except Exception:
        pass

    doc.save(args.output, garbage=3, deflate=True)
    doc.close()

    try:
        shutil.copystat(args.input, args.output)
    except Exception:
        pass

    if not args.quiet:
        print("\n本地 Paddle 双层完成:")
        print(f"  新增页面: {inserted_pages}/{total_pages}")
        print(f"  新增文本块: {inserted_blocks}")
        print(f"  跳过页面: {skipped_pages}")

    args.backend_used = "local_paddle_layered"