"""End-to-end tests for the QUE-221 rebuild of CamelotLLMTableExtractor.

The real Camelot + pdf2image stack is stubbed; the test focuses on the
pipeline glue (parallel flavors -> classifier -> unifier -> correction)
and on the new provenance fields landing on ExtractedTable.
"""

import asyncio
from pathlib import Path
from typing import List

import pytest
from _pytest.monkeypatch import MonkeyPatch

from quber.agents.classifier import ClassifierResult, MockClassifier
from quber.agents.llm_client import LLMTableCorrection, MockLLMClient
from quber.agents.unifier import MockUnifier, UnifiedTable, UnifierResult
from quber.core.extractors.camelot.acquire import CamelotCandidate, is_content_empty
from quber.core.extractors.camelot.llm import CamelotLLMTableExtractor, pipeline


def stub_pdf(tmp_path: Path) -> Path:
    pdf = tmp_path / "doc.pdf"
    pdf.write_bytes(b"%PDF-stub")
    return pdf


def stub_renders(tmp_path: Path, pages: int = 1) -> List[Path]:
    out: List[Path] = []
    for i in range(1, pages + 1):
        p = tmp_path / f"page-{i:04d}.png"
        p.write_bytes(b"\x89PNG\r\n\x1a\n")
        out.append(p)
    return out


def test_is_content_empty_recognizes_empty_shells():
    assert is_content_empty("")
    assert is_content_empty("|   |   |\n|---|---|\n|   |   |")
    assert is_content_empty("|  |  |\n|  |  |")
    assert not is_content_empty("| a | b |\n|---|---|\n| 1 | 2 |")
    assert not is_content_empty("| | x |\n|---|---|")


def test_pipeline_filters_empty_shells_before_classify(tmp_path: Path, monkeypatch: MonkeyPatch):
    pdf = stub_pdf(tmp_path)

    monkeypatch.setattr(pipeline, "render_pages", lambda *a, **kw: stub_renders(tmp_path, 1))

    empty_shell = CamelotCandidate(
        candidate_id="lattice-p1-0",
        flavor="lattice",
        page=1,
        accuracy=0.0,
        markdown="|   |   |\n|---|---|\n|   |   |",
    )
    real = CamelotCandidate(
        candidate_id="stream-p1-0",
        flavor="stream",
        page=1,
        accuracy=98.0,
        markdown="| col | val |\n|---|---|\n| Revenue | 100 |",
    )
    monkeypatch.setattr(
        pipeline,
        "run_camelot_flavors_parallel",
        lambda *a, **kw: [empty_shell, real],
    )

    extractor = CamelotLLMTableExtractor(
        llm_client=MockLLMClient(),
        classifier=MockClassifier(ClassifierResult(is_table=True, reason="mock-yes")),
        unifier=MockUnifier(),
    )
    out = asyncio.run(extractor.extract_tables(pdf))

    assert len(out) == 1
    assert out[0].flavor == "stream"
    assert out[0].camelot_accuracy == 98.0
    assert out[0].classifier_decision is not None
    assert out[0].classifier_decision.is_table is True


def test_pipeline_drops_classifier_negatives(tmp_path: Path, monkeypatch: MonkeyPatch):
    pdf = stub_pdf(tmp_path)
    monkeypatch.setattr(pipeline, "render_pages", lambda *a, **kw: stub_renders(tmp_path, 1))
    monkeypatch.setattr(
        pipeline,
        "run_camelot_flavors_parallel",
        lambda *a, **kw: [
            CamelotCandidate(
                candidate_id="lattice-p1-0",
                flavor="lattice",
                page=1,
                accuracy=42.0,
                markdown="| Note | (1) Some footnote text |",
            )
        ],
    )

    extractor = CamelotLLMTableExtractor(
        llm_client=MockLLMClient(),
        classifier=MockClassifier(ClassifierResult(is_table=False, reason="footnote")),
        unifier=MockUnifier(),
    )
    out = asyncio.run(extractor.extract_tables(pdf))
    assert out == []


def test_pipeline_records_llm_corrected_flag(tmp_path: Path, monkeypatch: MonkeyPatch):
    pdf = stub_pdf(tmp_path)
    monkeypatch.setattr(pipeline, "render_pages", lambda *a, **kw: stub_renders(tmp_path, 1))
    cand = CamelotCandidate(
        candidate_id="lattice-p1-0",
        flavor="lattice",
        page=1,
        accuracy=88.0,
        markdown="| a | b |\n|---|---|\n| 1 | 2 |",
    )
    monkeypatch.setattr(pipeline, "run_camelot_flavors_parallel", lambda *a, **kw: [cand])

    forced_correction = LLMTableCorrection(
        title="Sample",
        caption="",
        markdown="| a | b |\n|---|---|\n| 1 | 2 |\n| 3 | 4 |",
        footnotes=[],
        header_rows=1,
    )
    extractor = CamelotLLMTableExtractor(
        llm_client=MockLLMClient(correction=forced_correction),
        classifier=MockClassifier(),
        unifier=MockUnifier(),
    )
    out = asyncio.run(extractor.extract_tables(pdf))
    assert len(out) == 1
    assert out[0].llm_corrected is True
    assert out[0].title == "Sample"


def test_pipeline_no_correction_means_llm_corrected_false(tmp_path: Path, monkeypatch: MonkeyPatch):
    pdf = stub_pdf(tmp_path)
    monkeypatch.setattr(pipeline, "render_pages", lambda *a, **kw: stub_renders(tmp_path, 1))
    cand = CamelotCandidate(
        candidate_id="lattice-p1-0",
        flavor="lattice",
        page=1,
        accuracy=88.0,
        markdown="| a | b |\n|---|---|\n| 1 | 2 |",
    )
    monkeypatch.setattr(pipeline, "run_camelot_flavors_parallel", lambda *a, **kw: [cand])

    extractor = CamelotLLMTableExtractor(
        llm_client=MockLLMClient(),
        classifier=MockClassifier(),
        unifier=MockUnifier(),
    )
    out = asyncio.run(extractor.extract_tables(pdf))
    assert len(out) == 1
    assert out[0].llm_corrected is False


def test_pipeline_empty_camelot_returns_empty(tmp_path: Path, monkeypatch: MonkeyPatch):
    pdf = stub_pdf(tmp_path)
    monkeypatch.setattr(pipeline, "render_pages", lambda *a, **kw: stub_renders(tmp_path, 1))
    monkeypatch.setattr(pipeline, "run_camelot_flavors_parallel", lambda *a, **kw: [])

    extractor = CamelotLLMTableExtractor(
        llm_client=MockLLMClient(),
        classifier=MockClassifier(),
        unifier=MockUnifier(),
    )
    out = asyncio.run(extractor.extract_tables(pdf))
    assert out == []


def test_pipeline_output_ordered_by_page_then_row(tmp_path: Path, monkeypatch: MonkeyPatch):
    """The graph's correction fan-out completes in arbitrary order; the
    final step must restore pages-ascending, row-within-page order."""
    pdf = stub_pdf(tmp_path)
    monkeypatch.setattr(pipeline, "render_pages", lambda *a, **kw: stub_renders(tmp_path, 3))

    def cand(flavor: str, page: int, idx: int) -> CamelotCandidate:
        return CamelotCandidate(
            candidate_id=f"{flavor}-p{page}-{idx}",
            flavor=flavor,  # type: ignore[arg-type]
            page=page,
            accuracy=90.0,
            markdown=f"| col |\n|---|\n| p{page}-{idx} |",
        )

    # Deliberately out of page order.
    candidates = [cand("stream", 3, 0), cand("lattice", 1, 0), cand("stream", 2, 0)]
    monkeypatch.setattr(pipeline, "run_camelot_flavors_parallel", lambda *a, **kw: candidates)

    extractor = CamelotLLMTableExtractor(
        llm_client=MockLLMClient(),
        classifier=MockClassifier(),
        unifier=MockUnifier(),
    )
    out = asyncio.run(extractor.extract_tables(pdf))
    assert [t.page for t in out] == [1, 2, 3]


def test_pipeline_unifies_multi_candidate_pages(tmp_path: Path, monkeypatch: MonkeyPatch):
    """Two flavor outputs on one page flow through the unifier and come
    back as one canonical table crediting both candidates."""
    pdf = stub_pdf(tmp_path)
    monkeypatch.setattr(pipeline, "render_pages", lambda *a, **kw: stub_renders(tmp_path, 1))

    lattice = CamelotCandidate(
        candidate_id="lattice-p1-0",
        flavor="lattice",
        page=1,
        accuracy=95.0,
        markdown="| a | b |\n|---|---|\n| 1 | 2 |",
    )
    stream = CamelotCandidate(
        candidate_id="stream-p1-0",
        flavor="stream",
        page=1,
        accuracy=80.0,
        markdown="| a | b |\n|---|---|\n| 1 | 2 |\n| 3 | 4 |",
    )
    monkeypatch.setattr(pipeline, "run_camelot_flavors_parallel", lambda *a, **kw: [lattice, stream])

    unified_md = "| a | b |\n|---|---|\n| 1 | 2 |\n| 3 | 4 |"
    unifier = MockUnifier(
        UnifierResult(
            tables=[
                UnifiedTable(
                    markdown=unified_md,
                    source_candidate_ids=["lattice-p1-0", "stream-p1-0"],
                )
            ]
        )
    )
    extractor = CamelotLLMTableExtractor(
        llm_client=MockLLMClient(),
        classifier=MockClassifier(),
        unifier=unifier,
    )
    out = asyncio.run(extractor.extract_tables(pdf))
    assert len(out) == 1
    assert out[0].markdown == unified_md
    # Primary candidate is the first source id: lattice.
    assert out[0].flavor == "lattice"
    assert out[0].camelot_accuracy == 95.0


def test_pipeline_unifier_sees_camelot_order_despite_classify_completion_order(
    tmp_path: Path, monkeypatch: MonkeyPatch
):
    """The unifier LLM is order-sensitive: each page's candidates must
    arrive in Camelot's output order (lattice before stream) even when
    the classify fan-out completes in reverse (QUE-241 regression;
    observed digit-altering unify merges on shuffled input)."""
    pdf = stub_pdf(tmp_path)
    monkeypatch.setattr(pipeline, "render_pages", lambda *a, **kw: stub_renders(tmp_path, 1))

    lattice = CamelotCandidate(
        candidate_id="lattice-p1-0",
        flavor="lattice",
        page=1,
        accuracy=95.0,
        markdown="| a | b |\n|---|---|\n| 1 | 2 |",
    )
    stream = CamelotCandidate(
        candidate_id="stream-p1-0",
        flavor="stream",
        page=1,
        accuracy=80.0,
        markdown="| a | b |\n|---|---|\n| 1 | 2 |\n| 3 | 4 |",
    )
    monkeypatch.setattr(pipeline, "run_camelot_flavors_parallel", lambda *a, **kw: [lattice, stream])

    class SlowFirstClassifier:
        """Delays the lattice candidate so stream finishes classify first."""

        async def classify(self, markdown: str) -> ClassifierResult:
            if "| 3 | 4 |" not in markdown:  # the lattice candidate
                await asyncio.sleep(0.05)
            return ClassifierResult(is_table=True, reason="mock")

    class SpyUnifier(MockUnifier):
        def __init__(self) -> None:
            super().__init__()
            self.seen_orders: List[List[str]] = []

        async def unify(self, page_image, candidates):  # type: ignore[override]
            self.seen_orders.append([c.candidate_id for c in candidates])
            return await super().unify(page_image, candidates)

    spy = SpyUnifier()
    extractor = CamelotLLMTableExtractor(
        llm_client=MockLLMClient(),
        classifier=SlowFirstClassifier(),
        unifier=spy,
    )
    out = asyncio.run(extractor.extract_tables(pdf))
    assert len(out) == 2
    assert spy.seen_orders == [["lattice-p1-0", "stream-p1-0"]]


def test_pipeline_raises_on_missing_source(tmp_path: Path):
    extractor = CamelotLLMTableExtractor(
        llm_client=MockLLMClient(),
        classifier=MockClassifier(),
        unifier=MockUnifier(),
    )
    with pytest.raises(FileNotFoundError):
        asyncio.run(extractor.extract_tables(tmp_path / "missing.pdf"))