# tests/test_retriever.py
"""Tests for app.services.retriever.Retriever.hybrid_search — the BM25 + ColBERT merge."""
from __future__ import annotations

import sqlite3
from pathlib import Path
from unittest.mock import MagicMock

import pytest

from app.services.bm25_index import BM25Index
from app.services.retriever import Retriever


@pytest.fixture
def seeded_db(tmp_path) -> str:
    db_path = str(tmp_path / "test.db")
    schema = Path("migrations/001_initial_schema.sql").read_text()
    conn = sqlite3.connect(db_path)
    conn.executescript(schema)
    conn.execute(
        "INSERT INTO documents(id, title, file_path, status) VALUES ('d1','Test','test.pdf','ready')"
    )
    conn.execute(
        "INSERT INTO page_chunks(id, doc_id, page_number, text, source, word_count) "
        "VALUES ('c1','d1',1,'Setting the IP address on the AVC-X device','text',7)"
    )
    conn.execute(
        "INSERT INTO page_chunks(id, doc_id, page_number, text, source, word_count) "
        "VALUES ('c2','d1',2,'Filter cartridge replacement procedure','text',4)"
    )
    # Third, unrelated chunk — with only 2 chunks total, a term appearing in
    # exactly one of them gets an Okapi BM25 IDF of exactly log(1.0) == 0
    # (log((N-n+0.5)/(n+0.5)) with N=2, n=1), silently zeroing every score.
    # A third chunk dilutes N enough for real term-overlap scores to surface.
    conn.execute(
        "INSERT INTO page_chunks(id, doc_id, page_number, text, source, word_count) "
        "VALUES ('c3','d1',3,'Warranty terms and annual maintenance schedule','text',5)"
    )
    conn.commit()
    conn.close()
    return db_path


def _seeded_bm25() -> BM25Index:
    idx = BM25Index()
    idx._dirty = True
    return idx


def test_hybrid_search_falls_back_to_bm25_only_without_llm(seeded_db):
    retriever = Retriever(_seeded_bm25(), colbert=MagicMock())
    results = retriever.hybrid_search(
        query="IP address", top_k=5, doc_ids=None,
        db_path=seeded_db, vec_db_path="unused", llm=None,
    )
    assert any(r.chunk_id == "c1" for r in results)


def test_hybrid_search_falls_back_to_bm25_only_without_colbert(seeded_db):
    retriever = Retriever(_seeded_bm25(), colbert=None)
    results = retriever.hybrid_search(
        query="IP address", top_k=5, doc_ids=None,
        db_path=seeded_db, vec_db_path="unused", llm=MagicMock(),
    )
    assert any(r.chunk_id == "c1" for r in results)


def test_hybrid_search_merges_bm25_and_colbert_hits(seeded_db):
    fake_colbert = MagicMock()
    fake_colbert.query.return_value = [
        {"chunk_id": "c1", "doc_id": "d1", "page_number": 1, "text": "Setting the IP address on the AVC-X device", "score": 15.0},
        {"chunk_id": "c2", "doc_id": "d1", "page_number": 2, "text": "Filter cartridge replacement procedure", "score": 5.0},
    ]

    retriever = Retriever(_seeded_bm25(), colbert=fake_colbert)
    results = retriever.hybrid_search(
        query="IP address AVC-X", top_k=5, doc_ids=None,
        db_path=seeded_db, vec_db_path="unused", llm=MagicMock(),
    )

    fake_colbert.ensure_fresh.assert_called_once_with(seeded_db)
    result_ids = {r.chunk_id for r in results}
    assert "c1" in result_ids
    c1 = next(r for r in results if r.chunk_id == "c1")
    assert c1.vector_score == 1.0  # highest colbert score, normalized to max


def test_hybrid_search_falls_back_when_colbert_raises(seeded_db):
    fake_colbert = MagicMock()
    fake_colbert.query.side_effect = RuntimeError("model not loaded")

    retriever = Retriever(_seeded_bm25(), colbert=fake_colbert)
    results = retriever.hybrid_search(
        query="IP address", top_k=5, doc_ids=None,
        db_path=seeded_db, vec_db_path="unused", llm=MagicMock(),
    )

    assert any(r.chunk_id == "c1" for r in results)


def test_hybrid_search_discards_pure_noise(seeded_db):
    fake_colbert = MagicMock()
    fake_colbert.query.return_value = []

    retriever = Retriever(_seeded_bm25(), colbert=fake_colbert)
    results = retriever.hybrid_search(
        query="completely unrelated gibberish xyzzy",
        top_k=5, doc_ids=None,
        db_path=seeded_db, vec_db_path="unused", llm=MagicMock(),
    )

    assert results == []