"""Unit: retriever — ordering, dedup, and the context cap (fake rows).

The SQL side of :func:`app.rag.retriever.retrieve` is exercised by the
chat integration tests against real Postgres; the pure mapping logic in
:func:`select_documents` is tested here with in-memory rows.
"""
from __future__ import annotations

import uuid

from app.models import Document
from app.rag.retriever import TRUNCATION_MARKER, RetrievedChunk, select_documents


def _doc(path: str, content: str, source: str = "Homelab", title: str | None = None) -> Document:
    return Document(
        id=uuid.uuid4(),
        source=source,
        path=path,
        full_path=f"/tmp/{path}",
        title=title or path,
        content=content,
        content_hash="0" * 64,
    )


def _chunk(doc: Document, score: float, position: int = 0) -> RetrievedChunk:
    return RetrievedChunk(
        chunk_id=uuid.uuid4(),
        position=position,
        content=doc.content[:40],
        score=score,
        document=doc,
    )


def test_ranks_by_best_chunk_score_not_first_hit() -> None:
    """A doc whose *later* chunk scores highest must still rank first."""
    a = _doc("a.md", "A" * 50)
    b = _doc("b.md", "B" * 50)
    c = _doc("c.md", "C" * 50)
    chunks = [
        _chunk(a, 0.4, position=0),  # a's weak chunk comes first
        _chunk(b, 0.8),
        _chunk(a, 0.9, position=2),  # a's best chunk comes last
        _chunk(c, 0.5),
    ]
    docs = select_documents(chunks, n=3, max_chars=10_000)
    assert [d.path for d in docs] == ["a.md", "b.md", "c.md"]


def test_dedups_to_one_document_per_hit_set() -> None:
    a = _doc("a.md", "A" * 50)
    chunks = [_chunk(a, 0.2), _chunk(a, 0.7), _chunk(a, 0.5)]
    docs = select_documents(chunks, n=2, max_chars=10_000)
    assert len(docs) == 1
    assert docs[0] is a


def test_caps_at_n_documents() -> None:
    docs_in = [_doc(f"d{i}.md", "X" * 20) for i in range(4)]
    chunks = [_chunk(d, 0.5 - 0.1 * i) for i, d in enumerate(docs_in)]
    out = select_documents(chunks, n=2, max_chars=10_000)
    assert [d.path for d in out] == ["d0.md", "d1.md"]


def test_combined_content_capped_with_truncation_marker() -> None:
    big = _doc("big.md", "B" * 100)
    small = _doc("small.md", "S" * 100)
    chunks = [_chunk(big, 0.9), _chunk(small, 0.6)]
    out = select_documents(chunks, n=2, max_chars=150)
    # Best doc stays intact; the overflowing one is truncated in place.
    assert out[0].content == "B" * 100
    assert out[1].content.endswith(TRUNCATION_MARKER)
    assert out[1].content.startswith("S")
    assert len(out[0].content) + len(out[1].content) <= 150


def test_single_doc_over_budget_is_truncated_to_budget() -> None:
    big = _doc("big.md", "Z" * 200)
    out = select_documents([_chunk(big, 0.9)], n=2, max_chars=50)
    assert len(out[0].content) == 50
    assert out[0].content.endswith(TRUNCATION_MARKER)


def test_under_budget_no_truncation() -> None:
    a = _doc("a.md", "A" * 80)
    b = _doc("b.md", "B" * 60)
    out = select_documents([_chunk(b, 0.5), _chunk(a, 0.9)], n=2, max_chars=200)
    assert [d.path for d in out] == ["a.md", "b.md"]
    assert a.content == "A" * 80 and b.content == "B" * 60
    assert TRUNCATION_MARKER not in a.content + b.content


def test_empty_hits_yield_no_documents() -> None:
    assert select_documents([], n=2, max_chars=24_000) == []