# -*- coding: utf-8 -*-
"""tools/builtin_tools/pdf_reader.py 的单元测试。

用 PyMuPDF 现场生成含文本与内嵌图片的测试 PDF，验证文本结构化提取、
图片提取落地，以及异常分支（文件不存在 / 超过大小上限）。

需在装有 PyMuPDF 的 haocode 环境运行::

    python -m unittest discover tests
"""
import os
import sys
import tempfile
import unittest

# 让 tests 目录能 import 到项目根目录下的 tools 包
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

try:
    import pymupdf
except ImportError:  # 兼容旧版导入名
    import fitz as pymupdf  # type: ignore

from tools.builtin_tools.pdf_reader import (
    extract_pdf_text,
    extract_pdf_images,
)


def _make_png_bytes() -> bytes:
    """生成一张 8x8 红色小图的 PNG 字节流。"""
    pix = pymupdf.Pixmap(pymupdf.csRGB, pymupdf.IRect(0, 0, 8, 8), 0)
    pix.set_rect(pix.irect, (255, 0, 0))
    data = pix.tobytes("png")
    pix = None
    return data


def _build_sample_pdf(path: str) -> None:
    """造一个 2 页 PDF：第 1 页含文本 + 图片，第 2 页仅文本。"""
    doc = pymupdf.open()
    p1 = doc.new_page()
    p1.insert_text((72, 72), "Hello PDF page one")
    p1.insert_image(pymupdf.Rect(72, 100, 172, 200), stream=_make_png_bytes())
    p2 = doc.new_page()
    p2.insert_text((72, 72), "Second page text here")
    doc.save(path)
    doc.close()


class PdfReaderTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)
        self.dir = self.tmp.name
        self.pdf = os.path.join(self.dir, "sample.pdf")
        _build_sample_pdf(self.pdf)

    def test_extract_text_structure(self):
        text, pages = extract_pdf_text(self.pdf)
        self.assertEqual(pages, 2)
        self.assertIn("[第 1 页]", text)
        self.assertIn("[第 2 页]", text)
        self.assertIn("Hello PDF page one", text)
        self.assertIn("Second page text here", text)

    def test_extract_images(self):
        out_dir = os.path.join(self.dir, "imgs")
        imgs = extract_pdf_images(self.pdf, out_dir)
        self.assertGreaterEqual(len(imgs), 1)
        im = imgs[0]
        for key in ("page", "index", "abs_path", "mime", "size_kb", "width", "height"):
            self.assertIn(key, im)
        self.assertEqual(im["page"], 1)
        self.assertEqual(im["mime"], "image/png")
        self.assertTrue(os.path.isfile(im["abs_path"]))
        self.assertGreater(os.path.getsize(im["abs_path"]), 0)

    def test_missing_file(self):
        with self.assertRaises(ValueError):
            extract_pdf_text(os.path.join(self.dir, "nope.pdf"))

    def test_oversize(self):
        # 用一个极小的上限触发超大分支
        with self.assertRaises(ValueError):
            extract_pdf_text(self.pdf, max_bytes=10)


if __name__ == "__main__":
    unittest.main()