"""Integration against the development Postgres and its ingested corpus. A
mocked query proves the mock, not the SQL, so these run the real statements:
reference resolution over the groundings table, the fused retrieval window
(which embeds the query, so the local embedding model loads), and the cell
index gold verification depends on. Skipped when the database is unreachable
or holds no fusion document."""

from __future__ import annotations

import pytest

from quber.playground import db

pytestmark = pytest.mark.integration


@pytest.fixture(scope="module")
def fusion_doc():
    """(id, doc_key) of one ingested fusion document."""
    try:
        with db.connect() as conn:
            row = conn.execute(
                "SELECT id, doc_key FROM ade_playground.documents "
                "WHERE ade_version = 'quber-fusion' ORDER BY doc_key LIMIT 1"
            ).fetchone()
    except Exception as exc:
        pytest.skip(f"development Postgres unreachable: {exc}")
    if not row:
        pytest.skip("no fusion document ingested")
    return row


def test_reference_resolution_against_known_cited_ids(fusion_doc):
    from quber.playground.app import _resolve_refs

    doc_id, _ = fusion_doc
    with db.connect() as conn:
        row = conn.execute(
            """SELECT g.ref_id, g.page, g.status FROM ade_playground.groundings g
               WHERE g.document_id = %s AND g.ref_type = 'tableCell' AND g.status IS NOT NULL
               ORDER BY g.ref_id LIMIT 2""",
            (doc_id,),
        ).fetchall()
    assert row, "expected at least one grounded cell with provenance"
    cited = [r[0] for r in row]

    refs = _resolve_refs(doc_id, cited)
    assert [r.ref_id for r in refs] == cited  # citation order preserved
    first = refs[0]
    assert first.page == (row[0][1] or 0) + 1  # 1-based for the viewer
    assert first.status == row[0][2]
    assert first.text is not None

    # An id the document does not carry resolves to nothing, never to a guess.
    assert _resolve_refs(doc_id, ["no-such-id"]) == []


def test_cited_cell_suppresses_its_line_record(fusion_doc):
    from quber.playground.app import _resolve_refs

    doc_id, _ = fusion_doc
    with db.connect() as conn:
        row = conn.execute(
            """SELECT c.ref_id, l.ref_id
               FROM ade_playground.groundings c
               JOIN ade_playground.groundings l
                 ON l.document_id = c.document_id
                AND l.ref_type = 'line_item'
                AND l.ref_id = split_part(c.ref_id, '-', 1) || '-line-' || (c.position->>'row')
               WHERE c.document_id = %s AND c.ref_type = 'tableCell' AND c.bbox IS NOT NULL
               LIMIT 1""",
            (doc_id,),
        ).fetchone()
    if not row:
        pytest.skip("no cell with an enclosing line record")
    cell_id, line_id = row

    ids = [r.ref_id for r in _resolve_refs(doc_id, [line_id, cell_id])]
    assert cell_id in ids and line_id not in ids

    # The record keeps its box when no cell of its own is cited beside it.
    assert [r.ref_id for r in _resolve_refs(doc_id, [line_id])] == [line_id]


@pytest.mark.slow
def test_fused_window_over_the_real_corpus(fusion_doc):
    import asyncio

    from quber.playground.retrieval import _cap_line_records, _fused_window

    _, doc_key = fusion_doc
    window = asyncio.run(_fused_window(doc_key, "total revenue for the quarter", 30))
    assert window, "fused window returned nothing over an ingested document"
    scores = [c.score for c in window]
    assert scores == sorted(scores, reverse=True)
    assert len({c.chunk_id for c in window}) == len(window)
    capped = _cap_line_records(window, 2)
    per_parent = {}
    for c in capped:
        if c.chunk_type == "line_item" and c.parent_chunk_id:
            per_parent[c.parent_chunk_id] = per_parent.get(c.parent_chunk_id, 0) + 1
    assert all(n <= 2 for n in per_parent.values())


def test_cell_index_over_the_real_corpus(fusion_doc):
    from quber.playground.benchmark.verify_gold import cell_index, occurrences

    _, doc_key = fusion_doc
    index = cell_index(doc_key)
    assert index, "expected tagged cells in an ingested fusion document"
    (_, cid), text = next(iter(index.items()))
    # The chunk-less key resolves the same cell.
    assert index[("", cid)] is not None
    if text:
        assert occurrences(index, text) >= 1