"""Tests for app.ai.post_process — mirrors postProcess.test.ts."""
import pytest
from ai_citer.ai.post_process import assign_page_numbers
from ai_citer.models import Document, DocumentContent, TextSegment, Fact, Citation


# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------

def make_pdf_doc(segments: list[dict]) -> Document:
    segs = [TextSegment(**s) for s in segments]
    raw_text = " ".join(s["text"] for s in segments)
    return Document(
        id="test-doc",
        title="Test PDF",
        type="pdf",
        content=DocumentContent(rawText=raw_text, segments=segs),
        createdAt="2024-01-01T00:00:00Z",
    )


def make_text_doc() -> Document:
    return Document(
        id="text-doc",
        title="Test Text",
        type="text",
        content=DocumentContent(
            rawText="hello world",
            segments=[TextSegment(text="hello world", charOffset=0)],
        ),
        createdAt="2024-01-01T00:00:00Z",
    )


def make_fact(char_offset: int, length: int = 5) -> Fact:
    return Fact(
        id="f1",
        text="some fact",
        citations=[Citation(quoteText="quote", charOffset=char_offset, length=length, confidence="exact")],
    )


# ---------------------------------------------------------------------------
# non-PDF documents
# ---------------------------------------------------------------------------

class TestNonPdfDocuments:
    def test_does_nothing_for_text_documents(self):
        doc = make_text_doc()
        f = make_fact(0)
        assign_page_numbers(doc, [f])
        assert f.citations[0].pageNumber is None

    def test_does_nothing_when_segments_empty(self):
        doc = Document(
            id="x", title="x", type="pdf",
            content=DocumentContent(rawText="text", segments=[]),
            createdAt="2024-01-01T00:00:00Z",
        )
        f = make_fact(0)
        assign_page_numbers(doc, [f])
        assert f.citations[0].pageNumber is None


# ---------------------------------------------------------------------------
# PDF page number assignment
# ---------------------------------------------------------------------------

class TestPdfPageNumbers:
    def test_assigns_correct_page_number(self):
        doc = make_pdf_doc([
            {"text": "Page one content here.", "charOffset": 0, "pageNumber": 1},
            {"text": "Page two content here.", "charOffset": 30, "pageNumber": 2},
        ])
        f = make_fact(35)  # within page 2 segment (offset 30, length 22)
        assign_page_numbers(doc, [f])
        assert f.citations[0].pageNumber == 2

    def test_assigns_page_1_at_offset_zero(self):
        doc = make_pdf_doc([
            {"text": "First page text.", "charOffset": 0, "pageNumber": 1},
            {"text": "Second page text.", "charOffset": 20, "pageNumber": 2},
        ])
        f = make_fact(0)
        assign_page_numbers(doc, [f])
        assert f.citations[0].pageNumber == 1

    def test_skips_citations_with_minus_one_offset(self):
        doc = make_pdf_doc([
            {"text": "Some text here.", "charOffset": 0, "pageNumber": 1},
        ])
        f = Fact(
            id="f1", text="unmatched",
            citations=[Citation(quoteText="missing", charOffset=-1, length=7, confidence="fuzzy")],
        )
        assign_page_numbers(doc, [f])
        assert f.citations[0].pageNumber is None

    def test_multiple_citations_across_pages(self):
        doc = make_pdf_doc([
            {"text": "Alpha text on page one.", "charOffset": 0, "pageNumber": 1},
            {"text": "Beta text on page two.", "charOffset": 30, "pageNumber": 2},
        ])
        f = Fact(
            id="f1", text="multi-citation fact",
            citations=[
                Citation(quoteText="Alpha", charOffset=0, length=5, confidence="exact"),
                Citation(quoteText="Beta", charOffset=30, length=4, confidence="exact"),
            ],
        )
        assign_page_numbers(doc, [f])
        assert f.citations[0].pageNumber == 1
        assert f.citations[1].pageNumber == 2

    def test_multiple_facts_each_get_their_own_page(self):
        doc = make_pdf_doc([
            {"text": "Page one has this.", "charOffset": 0, "pageNumber": 1},
            {"text": "Page two has that.", "charOffset": 25, "pageNumber": 2},
        ])
        facts = [make_fact(0), make_fact(25)]
        assign_page_numbers(doc, facts)
        assert facts[0].citations[0].pageNumber == 1
        assert facts[1].citations[0].pageNumber == 2

    def test_leaves_page_number_none_when_between_segments(self):
        doc = make_pdf_doc([
            {"text": "Seg A", "charOffset": 0, "pageNumber": 1},
            {"text": "Seg B", "charOffset": 100, "pageNumber": 2},
        ])
        # charOffset 50 is between the two segments
        f = make_fact(50)
        assign_page_numbers(doc, [f])
        assert f.citations[0].pageNumber is None
