cad-agent/tests/test_sources_pdf.py

"""`dwg_review.sources` 的 **PDF** 部分单测。

用手工构造的最小 PDF（`_build_pdf`）驱动，不依赖任何 PDF 生成库、不依赖外部样本文件。
真实工程图纸的端到端验收见仓库根 `verify_pdf_source.py`。
"""

import base64

import pytest

from dwg_review import config, sources
from tests.fakes import make_ctx


def _build_pdf(pages):
    """构造最小可解析 PDF。pages = [{"mediabox", "rotate", "text"}, ...]。"""
    objs = [
        b"<< /Type /Catalog /Pages 2 0 R >>",
        None,  # obj 2：Pages（末尾回填，需先知道各页 obj 号）
        b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
    ]
    page_ids = []
    for spec in pages:
        mb = spec.get("mediabox", (0, 0, 200, 100))
        rot = spec.get("rotate", 0)
        text = spec.get("text", "HELLO")
        parts = []
        if spec.get("fill"):
            # 整页填黑 → 墨迹铺满整页（用于「整页一个簇」的场景）
            parts.append(
                f"{mb[0]} {mb[1]} {mb[2] - mb[0]} {mb[3] - mb[1]} re f"
            )
        if text:
            parts.append(f"BT /F1 12 Tf 20 40 Td ({text}) Tj ET")
        stream = "\n".join(parts).encode()
        page_id = len(objs) + 1
        content_id = page_id + 1
        page_ids.append(page_id)
        objs.append(
            (
                f"<< /Type /Page /Parent 2 0 R "
                f"/MediaBox [{mb[0]} {mb[1]} {mb[2]} {mb[3]}] /Rotate {rot} "
                f"/Resources << /Font << /F1 3 0 R >> >> /Contents {content_id} 0 R >>"
            ).encode()
        )
        objs.append(
            b"<< /Length " + str(len(stream)).encode() + b" >>\nstream\n"
            + stream + b"\nendstream"
        )
    kids = " ".join(f"{i} 0 R" for i in page_ids)
    objs[1] = f"<< /Type /Pages /Kids [{kids}] /Count {len(page_ids)} >>".encode()

    out = bytearray(b"%PDF-1.4\n")
    offsets = []
    for i, body in enumerate(objs, 1):
        offsets.append(len(out))
        out += f"{i} 0 obj\n".encode() + body + b"\nendobj\n"
    xref_pos = len(out)
    n = len(objs) + 1
    out += f"xref\n0 {n}\n".encode() + b"0000000000 65535 f \n"
    for off in offsets:
        out += f"{off:010d} 00000 n \n".encode()
    out += (
        f"trailer\n<< /Size {n} /Root 1 0 R >>\nstartxref\n{xref_pos}\n%%EOF\n"
    ).encode()
    return bytes(out)


@pytest.fixture
def pdf_path(tmp_path):
    """把 `_build_pdf(pages)` 落盘，返回路径。"""

    def _make(pages):
        path = tmp_path / "mini.pdf"
        path.write_bytes(_build_pdf(pages))
        return str(path)

    return _make


# ---------------------------------------------------------------------------
#  PdfReader
# ---------------------------------------------------------------------------
def test_pdf_reader_frames_basic(pdf_path):
    reader = sources.PdfReader(pdf_path([{}]))
    try:
        frames = reader.frames()
    finally:
        reader.close()
    assert len(frames) == 1
    assert frames[0]["bbox"] == [0.0, 0.0, 200.0, 100.0]
    assert frames[0]["source_type"] == "pdf"
    assert frames[0]["page_index"] == 0
    assert frames[0]["sheet_scale"] == 1.0


def test_pdf_reader_frames_multipage_keeps_order(pdf_path):
    reader = sources.PdfReader(
        pdf_path([{"mediabox": (0, 0, 100, 50)}, {"mediabox": (0, 0, 300, 400)}])
    )
    try:
        frames = reader.frames()
    finally:
        reader.close()
    assert [f["page_index"] for f in frames] == [0, 1]
    assert frames[1]["bbox"] == [0.0, 0.0, 300.0, 400.0]


def test_pdf_reader_rot90_swaps_page_size(pdf_path):
    """带 /Rotate 90 的页：世界尺寸应是 (MediaBox 高, MediaBox 宽)。"""
    reader = sources.PdfReader(pdf_path([{"mediabox": (0, 0, 200, 100), "rotate": 90}]))
    try:
        frames = reader.frames()
    finally:
        reader.close()
    assert frames[0]["bbox"] == [0.0, 0.0, 100.0, 200.0]


def test_pdf_reader_render_region_size(pdf_path):
    reader = sources.PdfReader(pdf_path([{}]))
    try:
        out = reader.render((0, 0, 100, 50), 2.0)
    finally:
        reader.close()
    assert out.size == (200, 100)


def test_pdf_reader_render_bad_box(pdf_path):
    reader = sources.PdfReader(pdf_path([{}]))
    try:
        assert reader.render((0, 0, 0, 0), 2.0) is None
    finally:
        reader.close()


def test_pdf_reader_text_items(pdf_path):
    reader = sources.PdfReader(pdf_path([{"text": "HELLO"}]))
    try:
        items = reader.text_items(0)
    finally:
        reader.close()
    assert len(items) == 1
    assert items[0]["content"] == "HELLO"
    x0, y0, x1, y1 = items[0]["bbox"]
    assert 0 <= x0 < x1 <= 200 and 0 <= y0 < y1 <= 100


def test_pdf_reader_text_items_min_chars(pdf_path):
    """短于 PDF_TEXT_MIN_CHARS 的片段被丢弃（滤噪）。"""
    reader = sources.PdfReader(pdf_path([{"text": "AB"}]))
    try:
        assert reader.text_items(0, min_chars=8) == []
        assert len(reader.text_items(0, min_chars=2)) == 1
    finally:
        reader.close()


def test_pdf_reader_text_items_respects_box_filter(pdf_path):
    reader = sources.PdfReader(pdf_path([{"text": "HELLO"}]))
    try:
        # 文本在 y≈40~49，给一个完全错开的窗口 → 过滤掉
        assert reader.text_items(0, box=(0, 80, 200, 100)) == []
    finally:
        reader.close()


def test_pdf_reader_text_items_rot90_maps_into_world(pdf_path):
    reader = sources.PdfReader(pdf_path([{"mediabox": (0, 0, 200, 100), "rotate": 90}]))
    try:
        items = reader.text_items(0)
    finally:
        reader.close()
    assert len(items) == 1
    x0, y0, x1, y1 = items[0]["bbox"]
    # 世界尺寸为 (100, 200)（旋转后），bbox 必须落在其中
    assert 0 <= x0 < x1 <= 100.0
    assert 0 <= y0 < y1 <= 200.0


def test_pdf_reader_empty_text_layer(pdf_path):
    reader = sources.PdfReader(pdf_path([{"text": ""}]))
    try:
        assert reader.text_items(0) == []
    finally:
        reader.close()


# ---------------------------------------------------------------------------
#  PDF 端到端（离线，plan + collect）
# ---------------------------------------------------------------------------
def test_pdf_plan_and_collect_end_to_end(pdf_path, monkeypatch):
    monkeypatch.setattr(config, "TILING_ENABLED", True)
    path = pdf_path([{"text": "HELLO"}])
    ctx = make_ctx(connect_cad=False, source="pdf", source_path=path)
    try:
        plan = sources.plan_file(ctx)
        assert plan["drawing_info"]["unit"] == "pt"
        assert plan["drawing_info"]["page_count"] == 1
        assert len(plan["frames"]) == 1
        ev = sources.collect(ctx, 1, plan["frames"][0])
    finally:
        sources.close_reader(ctx)

    assert ev["source_type"] == "pdf"
    assert ev["window"] == {"min_x": 0.0, "min_y": 0.0, "max_x": 200.0, "max_y": 100.0}
    assert ev["entities"] == {"total": 0, "items": []}
    assert ev["texts"]["total"] == 1
    assert ev["texts"]["items"][0]["content"] == "HELLO"
    assert ev["images"] and ev["images"][0]["label"] == "overall"
    assert base64.b64decode(ev["images"][0]["b64"]).startswith(b"\x89PNG")


def test_pdf_pdf_max_pages_caps_frames(pdf_path, monkeypatch):
    monkeypatch.setattr(config, "PDF_MAX_PAGES", 1)
    path = pdf_path([{}, {}])
    ctx = make_ctx(connect_cad=False, source="pdf", source_path=path)
    try:
        plan = sources.plan_file(ctx)
    finally:
        sources.close_reader(ctx)
    assert len(plan["frames"]) == 1
    assert plan["drawing_info"]["page_count"] == 2


def test_pdf_blank_page_has_only_overall(pdf_path, monkeypatch):
    """无墨迹的页：墨迹网格为空 → 只出整体图（不凭空造瓦片）。"""
    monkeypatch.setattr(config, "TILING_ENABLED", True)
    path = pdf_path([{"text": ""}])
    ctx = make_ctx(connect_cad=False, source="pdf", source_path=path)
    try:
        plan = sources.plan_file(ctx)
        ev = sources.collect(ctx, 1, plan["frames"][0])
    finally:
        sources.close_reader(ctx)
    assert [i["label"] for i in ev["images"]] == ["overall"]


def test_pdf_collect_multi_page_uses_own_page(pdf_path):
    """多页时每页取证必须渲染**自己的页**（页序不可错位）。"""
    path = pdf_path(
        [
            {"mediabox": (0, 0, 200, 100), "text": "FIRST"},
            {"mediabox": (0, 0, 400, 200), "text": "SECOND"},
        ]
    )
    ctx = make_ctx(connect_cad=False, source="pdf", source_path=path)
    try:
        plan = sources.plan_file(ctx)
        ev2 = sources.collect(ctx, 2, plan["frames"][1])
    finally:
        sources.close_reader(ctx)
    assert ev2["window"]["max_x"] == 400.0
    assert ev2["texts"]["items"][0]["content"] == "SECOND"


def test_pdf_dedupes_tile_equal_to_overall(pdf_path, monkeypatch):
    """整页即一个簇且瓦片分辨率不更高时，不重复出一张与整体图相同的图。"""
    monkeypatch.setattr(config, "TILING_ENABLED", True)
    monkeypatch.setattr(config, "PDF_OVERVIEW_MAX_ZOOM", 4.0)
    monkeypatch.setattr(config, "PDF_TILE_MAX_ZOOM", 4.0)
    path = pdf_path([{"fill": True}])  # 墨迹铺满整页 → 整页即一个簇
    ctx = make_ctx(connect_cad=False, source="pdf", source_path=path)
    try:
        plan = sources.plan_file(ctx)
        ev = sources.collect(ctx, 1, plan["frames"][0])
    finally:
        sources.close_reader(ctx)
    labels = [i["label"] for i in ev["images"]]
    assert labels.count("overall") == 1
    assert "tile-1" not in labels  # 与整体图完全重合 → 省掉


def test_pdf_keeps_tile_when_higher_zoom_than_overall(pdf_path, monkeypatch):
    """瓦片分辨率**高于**整体图时不能去重（否则丢掉细节）。"""
    monkeypatch.setattr(config, "TILING_ENABLED", True)
    monkeypatch.setattr(config, "PDF_OVERVIEW_MAX_ZOOM", 1.0)
    monkeypatch.setattr(config, "PDF_TILE_MAX_ZOOM", 6.0)
    path = pdf_path([{"fill": True}])
    ctx = make_ctx(connect_cad=False, source="pdf", source_path=path)
    try:
        plan = sources.plan_file(ctx)
        ev = sources.collect(ctx, 1, plan["frames"][0])
    finally:
        sources.close_reader(ctx)
    assert [i["label"] for i in ev["images"]] == ["overall", "tile-1"]