from __future__ import annotations

import json
from unittest.mock import MagicMock

import pymupdf

from app.services import pdf_image_extractor as mod
from app.services.layout_detector import LayoutBox


def test_process_page_extracts_doclayout_caption(tmp_path):
    images_dest = tmp_path / "images"
    images_dest.mkdir()
    manifest: dict[str, dict] = {}

    pix = MagicMock()
    pix.tobytes.return_value = b"jpeg"

    page = MagicMock()
    page.rect.width = 600
    page.get_pixmap.return_value = pix
    page.get_text.return_value = "Figure 1: Overall architecture.\n"

    doc = MagicMock()
    doc.__getitem__.return_value = page

    boxes = [
        LayoutBox(100, 100, 300, 300, "picture"),
        LayoutBox(95, 310, 320, 325, "figure_caption"),
    ]

    extracted = mod._process_page(
        doc,
        0,
        boxes,
        images_dest=images_dest,
        manifest=manifest,
        seen_labels=set(),
        arxiv_id="2401.00001",
    )

    assert extracted == 1
    info = manifest["figure_(p1-1).jpg"]
    assert info["caption_text"] == "Figure 1: Overall architecture."
    assert info["caption_source"] == "doclayout"
    assert info["caption_box"] == [95.0, 310.0, 320.0, 325.0]


def test_process_page_includes_caption_in_render(tmp_path):
    """渲染时把 caption 区域合并进同一张截图。"""
    images_dest = tmp_path / "images"
    images_dest.mkdir()
    manifest: dict[str, dict] = {}

    pix = MagicMock()
    pix.tobytes.return_value = b"jpeg"

    page = MagicMock()
    page.rect.width = 600
    page.get_pixmap.return_value = pix
    page.get_text.return_value = "Figure 1: Caption text.\n"

    doc = MagicMock()
    doc.__getitem__.return_value = page

    boxes = [
        LayoutBox(100, 100, 300, 300, "picture"),
        LayoutBox(95, 310, 320, 325, "figure_caption"),
    ]

    mod._process_page(
        doc,
        0,
        boxes,
        images_dest=images_dest,
        manifest=manifest,
        seen_labels=set(),
        arxiv_id="2401.00001",
    )

    # 内容 [100,100,300,300] ∪ caption [95,310,320,325]，各方向加 _REGION_PADDING=5
    # → Rect(90, 95, 325, 330)
    clip = page.get_pixmap.call_args.kwargs["clip"]
    assert clip == pymupdf.Rect(90, 95, 325, 330)


def test_label_images_preserves_doclayout_caption(tmp_path, monkeypatch):
    arxiv_id = "2401.00001"
    paper_root = tmp_path / arxiv_id
    images_dest = paper_root / "images"
    images_dest.mkdir(parents=True)
    (images_dest / "figure_(p1-1).jpg").write_bytes(b"jpeg")
    (images_dest / "manifest.json").write_text(
        json.dumps(
            {
                "figure_(p1-1).jpg": {
                    "page": 1,
                    "type": "figure",
                    "label": "Figure (p1-1)",
                    "box": [100, 100, 300, 300],
                    "caption_text": "Figure 1: PDF original caption.",
                    "caption_source": "doclayout",
                }
            }
        )
    )

    pdf_path = tmp_path / "paper.pdf"
    pdf_path.write_bytes(b"%PDF")
    monkeypatch.setattr(mod, "paper_dir", lambda _arxiv_id: paper_root)

    page = MagicMock()
    page.search_for.return_value = [pymupdf.Rect(120, 305, 180, 320)]

    fake_doc = MagicMock()
    fake_doc.page_count = 1
    fake_doc.__getitem__.return_value = page
    fake_doc.__enter__.return_value = fake_doc
    fake_doc.__exit__.return_value = False
    monkeypatch.setattr(mod.pymupdf, "open", lambda _path: fake_doc)

    labeled = mod.label_images_by_summary(
        arxiv_id,
        [{"id": "Figure 1", "caption": "Summary caption."}],
        pdf_path=pdf_path,
    )

    assert labeled == 1
    manifest = json.loads((images_dest / "manifest.json").read_text())
    info = manifest["figure_1.jpg"]
    assert info["caption_text"] == "Figure 1: PDF original caption."
    assert info["caption_source"] == "doclayout"
    assert info["summary_caption_text"] == "Summary caption."