# -*- coding: utf-8 -*-
"""
PDF 读取工具 (pdf_reader)
=========================

项目内第一个正式内置工具（builtin tool）。对外提供两项能力：

  1. ``extract_pdf_text``   —— 提取 PDF 全部文本，按页结构化为提示词友好格式
  2. ``extract_pdf_images`` —— 提取 PDF 内嵌的所有图片，逐张落地为 PNG

代码规范与调用模式（约定）
--------------------------
* 每个对外函数都带完整类型注解（参数 + 返回值）。
* docstring 采用 ``Args / Returns / Raises`` 分段，逐参数说明含义与单位。
* 将来的工具注册器（tools/registry.py，尚未实现）可通过
  ``inspect.signature`` + ``__doc__`` 自动反射出 OpenAI 格式的 Tool JSON Schema，
  因此本文件的函数签名与文档即"接口契约"，请保持稳定、描述清晰。

当前阶段说明
------------
* 本工具暂**不接入** agent 的自动工具调用链路（registry 未实现），
  由 UI 附件流程（``ui/views/main_window.py`` 的 PDF 处理）直接 import 调用。
* 参数与返回值已按"可被自动调用"的标准描述清楚，后续接线无需改动函数本身。

依赖与协议
----------
* 依赖 PyMuPDF（导入名 ``pymupdf``）。
* PyMuPDF 为 **AGPL-3.0 / Artifex 商业** 双许可：分发或通过网络提供服务时，
  需按 AGPL 开源本项目或购买商业许可。纯自用/内部不分发则无约束。
"""
import os
import uuid

try:
    import pymupdf  # PyMuPDF >= 1.24 的官方导入名
except ImportError:  # 兼容旧版导入名 fitz
    import fitz as pymupdf  # type: ignore


# 单个 PDF 的大小上限（字节）：超过则拒绝解析，避免内存与上下文爆炸
PDF_MAX_BYTES = 20 * 1024 * 1024  # 20 MB


def _open_doc(path: str, max_bytes: int) -> "pymupdf.Document":
    """内部辅助：做存在性/大小校验后打开 PDF，统一异常为 ValueError。"""
    if not os.path.isfile(path):
        raise ValueError("文件不存在")
    size = os.path.getsize(path)
    if size > max_bytes:
        raise ValueError(
            f"PDF 过大（{size / 1024 / 1024:.1f} MB > {max_bytes / 1024 / 1024:.0f} MB）"
        )
    try:
        doc = pymupdf.open(path)
    except Exception as e:  # pymupdf 抛出的异常类型不固定，统一兜底
        raise ValueError(f"PDF 打开失败（可能损坏）: {e}")
    if doc.needs_pass:
        doc.close()
        raise ValueError("PDF 已加密，需要密码，暂不支持")
    return doc


def extract_pdf_text(path: str, max_bytes: int = PDF_MAX_BYTES) -> tuple:
    """提取 PDF 的全部文本，按页结构化。

    输出格式（无文本的空页会被跳过）::

        [第 1 页]
        <该页文本>

        [第 2 页]
        <该页文本>

    Args:
        path: PDF 文件路径（绝对路径，或相对当前工作目录的路径）。
        max_bytes: 允许解析的最大文件字节数，默认 ``PDF_MAX_BYTES``（20MB）。

    Returns:
        ``(structured_text, page_count)`` 二元组：

        * ``structured_text`` (str): 按页拼接的结构化文本，各页之间以空行分隔；
          若整份 PDF 无文本则为空字符串 ``""``。
        * ``page_count`` (int): PDF 的总页数（含空页）。

    Raises:
        ValueError: 文件不存在 / 超过 ``max_bytes`` / PDF 损坏 / PDF 加密。
    """
    doc = _open_doc(path, max_bytes)
    try:
        page_count = doc.page_count
        parts = []
        for i, page in enumerate(doc, 1):
            text = page.get_text("text").strip()
            if text:
                parts.append(f"[第 {i} 页]\n{text}")
        return "\n\n".join(parts), page_count
    finally:
        doc.close()


def extract_pdf_images(path: str, out_dir: str, max_bytes: int = PDF_MAX_BYTES) -> list:
    """提取 PDF 内嵌的所有图片，逐张保存为 PNG。

    遍历每一页的内嵌图片（``page.get_images``），用 ``Pixmap`` 解码后写出；
    CMYK 色彩空间会自动转换为 RGB（否则无法存为 PNG）。单张图片解码失败
    会被跳过，不影响其余图片。

    Args:
        path: PDF 文件路径。
        out_dir: 图片输出目录（建议使用绝对路径）。不存在时自动创建。
        max_bytes: 允许解析的最大文件字节数，默认 ``PDF_MAX_BYTES``（20MB）。

    Returns:
        图片元数据列表（按 ``(page, index)`` 升序），每项为 dict::

            {
                "page":     int,    # 所在页码（从 1 开始）
                "index":    int,    # 该页内的第几张（从 1 开始）
                "abs_path": str,    # 落地 PNG 的绝对路径
                "mime":     str,    # 固定为 "image/png"
                "size_kb":  float,  # 文件大小（KB，保留两位小数）
                "width":    int,    # 图片像素宽
                "height":   int,    # 图片像素高
            }

        PDF 中无内嵌图片时返回空列表 ``[]``。

    Raises:
        ValueError: 文件不存在 / 超过 ``max_bytes`` / PDF 损坏 / PDF 加密。
    """
    doc = _open_doc(path, max_bytes)
    os.makedirs(out_dir, exist_ok=True)
    tag = uuid.uuid4().hex[:8]  # 本批次文件名前缀，避免跨 PDF 撞名
    results = []
    try:
        for page_num, page in enumerate(doc, 1):
            for img_index, img_info in enumerate(page.get_images(full=True), 1):
                xref = img_info[0]
                try:
                    pix = pymupdf.Pixmap(doc, xref)
                    if pix.n - pix.alpha >= 4:  # CMYK → RGB
                        pix = pymupdf.Pixmap(pymupdf.csRGB, pix)
                    fname = f"pdfimg_{tag}_p{page_num}_{img_index}.png"
                    abs_path = os.path.join(out_dir, fname)
                    pix.save(abs_path)
                    width, height = pix.width, pix.height
                    pix = None  # 及时释放位图内存
                except Exception:
                    continue  # 单张失败不影响整体
                results.append({
                    "page": page_num,
                    "index": img_index,
                    "abs_path": abs_path,
                    "mime": "image/png",
                    "size_kb": round(os.path.getsize(abs_path) / 1024, 2),
                    "width": width,
                    "height": height,
                })
        return results
    finally:
        doc.close()