tests/test_ai_highlight_engine.py

"""
Unit tests for the AI highlight engine (services/ai_highlight_engine.py):
sentence chunking, line->chunk alignment, colour mapping, prose-only
filtering, paragraph badge merging, figure-region suppression and the
full highlight + legend page pipeline.
"""
import fitz

from services.ai_highlight_engine import (
    chunk_text,
    align_lines_to_chunks,
    p_fake_to_color,
    highlight_original_pdf,
    _detect_visual_regions,
    _merge_rects,
    _is_analyzed,
    GREEN,
    AMBER,
    RED,
    BADGE_WIDTH,
    BADGE_HEIGHT,
    CHIP_OPACITY,
)


def test_chunk_text_groups_sentences_above_min_words():
    text = "First sentence here. " * 30
    chunks = chunk_text(text, min_words=20, max_words=110)
    assert chunks
    assert all(len(c.split()) >= 20 for c in chunks)


def test_chunk_text_preserves_every_word():
    text = ("Alpha sentence one. Beta sentence two. " * 40).strip()
    chunks = chunk_text(text, min_words=20, max_words=110)
    joined = " ".join(chunks)
    assert len(joined.split()) == len(text.split())
    assert joined == text


def test_chunk_text_empty_returns_empty():
    assert chunk_text("") == []
    assert chunk_text(None) == []


def test_chunk_text_merges_tiny_tail():
    text = "This is a long first chunk with enough words to be scored reliably here. " + "Short tail. " * 3
    chunks = chunk_text(text, min_words=20, max_words=110)
    assert chunks
    assert len(chunks[-1].split()) >= 20  # tail merged, not left tiny


def test_align_lines_to_chunks_in_order():
    chunks = [
        "The quick brown fox jumps over the lazy dog and keeps running fast.",
        "The hedgehog sleeps under the autumn leaves all through the winter.",
    ]
    lines = [
        "The quick brown fox jumps over the lazy dog",
        "and keeps running fast.",
        "The hedgehog sleeps under the autumn leaves",
        "all through the winter.",
    ]
    assignments = align_lines_to_chunks(lines, chunks)
    assert assignments == [0, 0, 1, 1]


def test_align_unmatched_line_keeps_nearest_chunk():
    chunks = [
        "The quick brown fox jumps over the lazy dog and keeps running fast.",
        "The hedgehog sleeps under the autumn leaves all through the winter.",
    ]
    lines = ["A line of text that matches nothing anywhere at all.", "The hedgehog sleeps under the autumn leaves"]
    assignments = align_lines_to_chunks(lines, chunks)
    assert len(assignments) == 2
    assert assignments[0] == 0  # nearest chunk, never dropped
    assert assignments[1] == 1


def test_p_fake_to_color_green_to_red():
    c0 = p_fake_to_color(0.0)
    c50 = p_fake_to_color(0.5)
    c100 = p_fake_to_color(1.0)
    assert c0 == GREEN
    assert c50 == AMBER
    assert c100 == RED
    # green dominates at low p_fake, red at high
    assert c0[1] > c0[0] and c0[1] > c0[2]
    assert c100[0] > c100[1] and c100[0] > c100[2]
    assert p_fake_to_color(-1) == GREEN
    assert p_fake_to_color(2) == RED


def _make_pdf(path, line_groups):
    doc = fitz.open()
    for group in line_groups:
        page = doc.new_page(width=595, height=842)
        for i, line in enumerate(group):
            page.insert_text((72, 100 + i * 40), line)
    doc.save(path)
    doc.close()


def test_highlight_original_pdf_annotates_every_line(tmp_path):
    pdf = tmp_path / "paper.pdf"
    _make_pdf(
        pdf,
        [
            ("The quick brown fox jumps over the lazy dog and keeps running fast.",
             "The hedgehog sleeps under the autumn leaves all through the winter."),
            ("Final page line one about neural networks and deep learning models.",
             "Final page line two concludes the paper with a clear summary."),
        ],
    )
    chunks = [
        {"chunk_id": "p1#0", "text": "The quick brown fox jumps over the lazy dog and keeps running fast. "
                                     "The hedgehog sleeps under the autumn leaves all through the winter.",
         "p_fake": 0.12, "verdict": "REAL", "confidence": 0.76},
        {"chunk_id": "p2#0", "text": "Final page line one about neural networks and deep learning models. "
                                     "Final page line two concludes the paper with a clear summary.",
         "p_fake": 0.88, "verdict": "AI_GENERATED", "confidence": 0.76},
    ]
    out = tmp_path / "out.pdf"
    stats = highlight_original_pdf(
        str(pdf), chunks, str(out),
        summary={"p_fake": 0.5, "verdict": "UNCERTAIN", "confidence": 0.1},
    )

    assert stats["total_lines"] == 4
    assert stats["analyzed_lines"] == 4
    assert stats["lines_highlighted"] == 4  # every analyzed prose line
    assert stats["pages_touched"] == 2
    assert "legend_page" not in stats  # legend page removed from the annex

    doc = fitz.open(str(out))
    assert doc.page_count == 2  # no legend page — just the 2 paper pages

    annots = [doc[p].annots() for p in range(2)]
    counts = [len(list(a)) if a is not None else 0 for a in annots]
    assert counts == [2, 2]  # every prose line highlighted

    # badges drawn with the verdict stated inline, not just a bare percent
    page1_text = doc[0].get_text()
    assert "HUMAN" in page1_text
    assert "12%" in page1_text
    page2_text = doc[1].get_text()
    assert "AI" in page2_text
    assert "88%" in page2_text
    doc.close()


def test_highlight_skips_page_furniture(tmp_path):
    pdf = tmp_path / "paper.pdf"
    doc = fitz.open()
    page = doc.new_page(width=595, height=842)
    page.insert_textbox(
        fitz.Rect(72, 80, 540, 200),
        "The quick brown fox jumps over the lazy dog. "
        "The hedgehog sleeps all through the winter night. "
        "The morning sun rises over the calm blue sea.",
        fontsize=11,
    )
    page.insert_text((72, 240), "Page 1 of 12", fontsize=9)
    page.insert_text((72, 260), "Figure 1. Results of the main experiment", fontsize=9)
    page.insert_text((72, 280), "References", fontsize=9)
    page.insert_text((72, 300), "Alice Johnson, Stanford University", fontsize=9)
    doc.save(str(pdf))
    doc.close()

    chunks = [
        {"chunk_id": "p#0",
         "text": "The quick brown fox jumps over the lazy dog. "
                 "The hedgehog sleeps all through the winter night. "
                 "The morning sun rises over the calm blue sea.",
         "p_fake": 0.2, "verdict": "REAL", "confidence": 0.7},
    ]
    out = tmp_path / "out.pdf"
    stats = highlight_original_pdf(str(pdf), chunks, str(out))

    assert stats["total_lines"] == 6
    assert stats["analyzed_lines"] == 2  # only the prose block
    assert stats["lines_highlighted"] == 2
    assert stats["badges"] == 1  # prose merged into a single paragraph badge

    doc = fitz.open(str(out))
    annots = list(doc[0].annots()) if doc[0].annots() else []
    assert len(annots) == 2  # furniture lines carry no highlight
    doc.close()


def test_highlight_merges_same_verdict_lines_into_one_badge(tmp_path):
    pdf = tmp_path / "paper.pdf"
    doc = fitz.open()
    page = doc.new_page(width=595, height=842)
    page.insert_textbox(
        fitz.Rect(72, 80, 540, 300),
        "The quick brown fox jumps over the lazy dog and keeps running fast. "
        "The hedgehog sleeps all through the winter night. "
        "The morning sun rises over the calm blue sea.",
        fontsize=11,
    )
    doc.save(str(pdf))
    doc.close()

    chunks = [
        {"chunk_id": "p#0",
         "text": "The quick brown fox jumps over the lazy dog and keeps running fast. "
                 "The hedgehog sleeps all through the winter night. "
                 "The morning sun rises over the calm blue sea.",
         "p_fake": 0.12, "verdict": "REAL", "confidence": 0.8},
    ]
    out = tmp_path / "out.pdf"
    stats = highlight_original_pdf(str(pdf), chunks, str(out))

    assert stats["lines_highlighted"] == 2
    assert stats["badges"] == 1  # one badge for the whole paragraph

    doc = fitz.open(str(out))
    assert doc[0].get_text().count("HUMAN 12%") == 1  # one badge, not one per line
    doc.close()


def test_merge_rects_unions_overlapping_into_one():
    """Regression: overlapping drawings must merge (PyMuPDF Rect has no
    .union(); the in-place union is include_rect). This crashed in
    production on any page with two overlapping visual regions."""
    rects = [
        fitz.Rect(72, 120, 400, 320),
        fitz.Rect(350, 280, 520, 420),
        fitz.Rect(500, 40, 560, 80),  # small, separate
    ]
    merged = _merge_rects(rects)
    assert len(merged) == 2  # two overlapping fused, one standalone
    fused = [r for r in merged if r.width > 400]
    assert fused and fused[0].x0 <= 72 and fused[0].y1 >= 420


def test_detect_visual_regions_merges_overlapping_drawings(tmp_path):
    pdf = tmp_path / "paper.pdf"
    doc = fitz.open()
    page = doc.new_page(width=595, height=842)
    page.draw_rect(fitz.Rect(72, 120, 400, 320), fill=(0.9, 0.9, 0.9), color=None)
    page.draw_rect(fitz.Rect(350, 280, 520, 420), fill=(0.8, 0.8, 0.8), color=None)
    page.insert_text((90, 150), "Model architecture overview diagram", fontsize=10)
    page.insert_text((90, 180), "encoder layer one output", fontsize=10)
    doc.save(str(pdf))
    regions = _detect_visual_regions(doc[0])
    assert len(regions) == 1  # the two overlapping diagram boxes fused
    assert regions[0].x0 <= 72 and regions[0].y1 >= 420
    doc.close()


def test_highlight_skips_lines_inside_figure_region(tmp_path):
    pdf = tmp_path / "paper.pdf"
    doc = fitz.open()
    page = doc.new_page(width=595, height=842)
    # diagram box (vector drawing -> detected visual region)
    page.draw_rect(fitz.Rect(72, 120, 400, 320), fill=(0.9, 0.9, 0.9), color=None)
    # labels inside the diagram
    page.insert_text((90, 150), "Model architecture overview diagram", fontsize=10)
    page.insert_text((90, 180), "encoder layer one output", fontsize=10)
    # prose outside the diagram
    page.insert_textbox(
        fitz.Rect(72, 360, 560, 500),
        "The proposed method improves accuracy on every benchmark dataset.",
        fontsize=11,
    )
    doc.save(str(pdf))
    regions = _detect_visual_regions(doc[0])
    assert len(regions) == 1  # region is detected
    doc.close()

    chunks = [
        {"chunk_id": "p#0",
         "text": "The proposed method improves accuracy on every benchmark dataset.",
         "p_fake": 0.15, "verdict": "REAL", "confidence": 0.8},
    ]
    out = tmp_path / "out.pdf"
    stats = highlight_original_pdf(str(pdf), chunks, str(out))

    assert stats["lines_highlighted"] == 1  # only the prose line survives
    assert stats["badges"] == 1

    doc2 = fitz.open(str(out))
    annots = list(doc2[0].annots()) if doc2[0].annots() else []
    assert len(annots) == 1
    doc2.close()


def test_badges_have_fixed_geometry_across_pages(tmp_path):
    pdf = tmp_path / "paper.pdf"
    doc = fitz.open()
    for i in range(2):
        page = doc.new_page(width=595, height=842)
        page.insert_text(
            (72, 100 + i * 400),
            f"The quick brown fox jumps over the lazy dog and runs far page {i}.",
            fontsize=14 if i == 0 else 8,  # different line heights per page
        )
    doc.save(str(pdf))
    doc.close()

    chunks = [
        {"chunk_id": f"p{i}#0",
         "text": f"The quick brown fox jumps over the lazy dog and runs far page {i}.",
         "p_fake": 0.3 + 0.4 * i, "verdict": "REAL", "confidence": 0.8}
        for i in range(2)
    ]
    out = tmp_path / "out.pdf"
    stats = highlight_original_pdf(str(pdf), chunks, str(out))
    assert stats["badges"] == 2

    doc2 = fitz.open(str(out))
    for pn in (0, 1):
        for d in doc2[pn].get_drawings():
            if abs((d.get("fill_opacity") or 0) - CHIP_OPACITY) < 0.01:
                assert abs(d["rect"].width - BADGE_WIDTH) < 0.1
                assert abs(d["rect"].height - BADGE_HEIGHT) < 0.1
    doc2.close()


def test_is_analyzed_rejects_furniture_and_short_lines():
    corpus = " ".join([
        "the quick brown fox jumps over the lazy dog",
        "the hedgehog sleeps all through the winter night",
    ])
    assert _is_analyzed("the quick brown fox jumps over the lazy dog", corpus)
    assert not _is_analyzed("page 1 of 12", corpus)          # header, not in corpus
    assert not _is_analyzed("references", corpus)            # too short
    assert not _is_analyzed("", corpus)


def test_highlight_original_pdf_empty_chunks_raises(tmp_path):
    pdf = tmp_path / "paper.pdf"
    _make_pdf(pdf, [("Some line one.", "Some line two.")])
    try:
        highlight_original_pdf(str(pdf), [], str(tmp_path / "out.pdf"))
        assert False, "expected ValueError"
    except ValueError:
        pass