"""Tests for hybrid search: explain traces, document_id, and RRF merging."""

import sys
import types

import pytest

from kb.database import (
    get_connection,
    init_schema,
    insert_chunk,
    insert_document,
    insert_embedding,
    tag_document,
)
from kb.search import _blend_rerank, _rank_by_score, _rrf_merge, hybrid_search

DIM = 4

# Chunk vectors are axis-aligned so we can steer vector ranking exactly:
# a query of [1,0,0,0] has distance 0 to chunk A, sqrt(2) to chunk B.
VEC_A = [1.0, 0.0, 0.0, 0.0]
VEC_B = [0.0, 1.0, 0.0, 0.0]
QUERY_VEC = [1.0, 0.0, 0.0, 0.0]


class _Cfg:
    search_threshold = 0.0
    rerank_enabled = False
    reranker_model = "test-reranker"
    rerank_candidates = 40


class _Db:
    def __init__(self, conn, ids):
        self.conn = conn
        self.ids = ids


@pytest.fixture
def db(tmp_path, monkeypatch):
    """Real schema (FTS5 + sqlite-vec) with two docs and stubbed embeddings."""
    fake = types.ModuleType("kb.embeddings")
    fake.embed_texts = lambda texts: [QUERY_VEC for _ in texts]
    monkeypatch.setitem(sys.modules, "kb.embeddings", fake)

    conn = get_connection(str(tmp_path / "kb.db"))
    init_schema(conn, embedding_dim=DIM)

    doc_a = insert_document(conn, "Alpha doc", "/src/a.md", "hash-a", "markdown")
    doc_b = insert_document(conn, "Bravo doc", "/src/b.md", "hash-b", "markdown")
    chunk_a = insert_chunk(conn, doc_a, 0, "alpha network switch configuration")
    chunk_b = insert_chunk(conn, doc_b, 0, "bravo unrelated cooking recipe")
    insert_embedding(conn, chunk_a, VEC_A)
    insert_embedding(conn, chunk_b, VEC_B)
    tag_document(conn, doc_a, ["ops"])

    ids = {"doc_a": doc_a, "doc_b": doc_b, "chunk_a": chunk_a, "chunk_b": chunk_b}
    yield _Db(conn, ids)
    conn.close()


def test_results_include_document_id(db):
    result = hybrid_search(db.conn, "alpha switch", _Cfg())
    assert result["results"], "expected at least one hit"
    top_hit = result["results"][0]
    assert top_hit["document_id"] == db.ids["doc_a"]
    assert top_hit["chunk_id"] == db.ids["chunk_a"]


def test_explain_absent_by_default(db):
    result = hybrid_search(db.conn, "alpha switch", _Cfg())
    assert all("explain" not in r for r in result["results"])


def test_explain_hybrid_breakdown(db):
    result = hybrid_search(db.conn, "alpha switch", _Cfg(), explain=True)
    top_hit = result["results"][0]
    exp = top_hit["explain"]

    # Chunk A is rank 1 in both arms: FTS matches "alpha"/"switch", vector
    # distance is 0 (similarity 1.0).
    assert exp["fts_rank"] == 1
    assert exp["vec_rank"] == 1
    assert exp["fts_score"] > 0
    assert exp["vec_score"] == pytest.approx(1.0)
    assert exp["rrf_fts"] == pytest.approx(1.0 / 61, abs=1e-6)
    assert exp["rrf_vec"] == pytest.approx(1.0 / 61, abs=1e-6)
    assert exp["final_score"] == pytest.approx(
        exp["rrf_fts"] + exp["rrf_vec"] + exp["bonus"], abs=1e-5
    )
    assert exp["final_score"] == pytest.approx(top_hit["score"], abs=1e-5)


def test_explain_single_arm_when_vec_misses(db):
    """A chunk found only by vector search has null FTS fields."""
    result = hybrid_search(db.conn, "zzz-no-fts-match", _Cfg(), explain=True)
    for r in result["results"]:
        exp = r["explain"]
        assert exp["fts_score"] is None
        assert exp["fts_rank"] is None
        assert exp["rrf_fts"] is None
        assert exp["vec_rank"] is not None


def test_explain_fts_only_shape(db):
    result = hybrid_search(db.conn, "alpha switch", _Cfg(), fts_only=True, explain=True)
    top_hit = result["results"][0]
    exp = top_hit["explain"]
    assert exp["fts_rank"] == 1
    assert exp["final_score"] == pytest.approx(exp["fts_score"])
    assert "vec_score" not in exp


def test_explain_vec_only_shape(db):
    result = hybrid_search(db.conn, "anything", _Cfg(), vec_only=True, explain=True)
    top_hit = result["results"][0]
    exp = top_hit["explain"]
    assert exp["vec_rank"] == 1
    assert exp["final_score"] == pytest.approx(exp["vec_score"])
    assert "fts_score" not in exp


def test_rrf_merge_arithmetic():
    fts = {1: 10.0, 2: 5.0}
    vec = {2: 0.9, 3: 0.8}
    scores, details = _rrf_merge(fts, vec)
    by_id = dict(scores)

    # fts rank 1 → 1/61 + 0.05 bonus
    assert by_id[1] == pytest.approx(1 / 61 + 0.05)
    # fts rank 2 (+0.02) and vec rank 1 (+0.05) — bonuses stack across arms
    assert by_id[2] == pytest.approx(1 / 62 + 1 / 61 + 0.07)
    # vec rank 2 → 1/62 + 0.02
    assert by_id[3] == pytest.approx(1 / 62 + 0.02)
    # Chunk 2 appears in both arms, so it must win.
    assert scores[0][0] == 2

    assert details[2]["fts_rank"] == 2
    assert details[2]["vec_rank"] == 1
    assert details[2]["bonus"] == pytest.approx(0.07)
    assert details[1]["vec_rank"] is None
    assert details[3]["rrf_fts"] is None


def test_top_rank_bonus_tiers():
    from kb.search import _top_rank_bonus

    assert _top_rank_bonus(1) == 0.05
    assert _top_rank_bonus(2) == 0.02
    assert _top_rank_bonus(3) == 0.02
    assert _top_rank_bonus(4) == 0.0
    assert _top_rank_bonus(None) == 0.0


def test_bonus_preserves_top_of_arm():
    """A chunk at rank 1 of one arm beats a chunk at mid-rank in both arms."""
    fts = {10: 100.0, 11: 90.0, 12: 80.0, 13: 70.0, 14: 60.0}
    vec = {20: 0.9, 11: 0.8, 12: 0.7, 13: 0.6, 14: 0.5}
    scores, _ = _rrf_merge(fts, vec)
    order = [cid for cid, _ in scores]
    # Chunk 10 (fts #1, absent from vec): 1/61 + 0.05 ≈ 0.0664.
    # Chunk 13 (rank 4 in fts, rank 4 in vec): 2/64 ≈ 0.031, no bonus.
    assert order.index(10) < order.index(13)


def test_rank_by_score():
    assert _rank_by_score({7: 0.5, 8: 0.9, 9: 0.1}) == {8: 1, 7: 2, 9: 3}


# ---------------------------------------------------------------------------
# Reranking
# ---------------------------------------------------------------------------

def _enable_fake_reranker(monkeypatch, score_fn):
    from kb import reranker

    monkeypatch.setattr(reranker, "is_available", lambda: True)
    monkeypatch.setattr(reranker, "rerank_scores", score_fn)


def test_rerank_flags_response_and_explain(db, monkeypatch):
    _enable_fake_reranker(monkeypatch, lambda q, texts: [0.9] * len(texts))
    result = hybrid_search(db.conn, "alpha switch", _Cfg(), explain=True, rerank=True)

    assert result["reranked"] is True
    top_hit = result["results"][0]
    exp = top_hit["explain"]
    assert exp["rerank_score"] == pytest.approx(0.9)
    assert exp["pre_rerank_rank"] == 1
    assert exp["blend_weight"] == 0.75
    assert exp["final_score"] == pytest.approx(top_hit["score"], abs=1e-5)
    # rank 1: 75% retrieval (norm=1.0 for the top candidate) + 25% rerank
    assert top_hit["score"] == pytest.approx(0.75 * 1.0 + 0.25 * 0.9, abs=1e-5)


def test_rerank_can_reorder(db, monkeypatch):
    """A strong rerank score rescues a lower-retrieval-ranked chunk."""

    def favour_chunk_b(query, texts):
        return [1.0 if "cooking" in t else 0.0 for t in texts]

    _enable_fake_reranker(monkeypatch, favour_chunk_b)
    # Neutral-ish query: both chunks retrieved, chunk A ranked first.
    result = hybrid_search(db.conn, "alpha cooking", _Cfg(), rerank=True)
    assert result["reranked"] is True
    assert len(result["results"]) == 2


def test_rerank_false_bypasses(db, monkeypatch):
    called = []
    _enable_fake_reranker(monkeypatch, lambda q, t: called.append(1) or [0.5] * len(t))

    result = hybrid_search(db.conn, "alpha switch", _Cfg(), rerank=False)
    assert result["reranked"] is False
    assert not called


def test_rerank_unavailable_degrades_gracefully(db):
    # No reranker loaded: rerank=True must not error.
    result = hybrid_search(db.conn, "alpha switch", _Cfg(), rerank=True, explain=True)
    assert result["reranked"] is False
    assert "rerank_score" not in result["results"][0]["explain"]


def test_rerank_skipped_for_single_arm(db, monkeypatch):
    called = []
    _enable_fake_reranker(monkeypatch, lambda q, t: called.append(1) or [0.5] * len(t))

    result = hybrid_search(db.conn, "alpha switch", _Cfg(), fts_only=True, rerank=True)
    assert result["reranked"] is False
    assert not called


def test_rerank_default_follows_cfg(db, monkeypatch):
    _enable_fake_reranker(monkeypatch, lambda q, texts: [0.5] * len(texts))

    class _RerankCfg(_Cfg):
        rerank_enabled = True

    result = hybrid_search(db.conn, "alpha switch", _RerankCfg())
    assert result["reranked"] is True


def test_blend_rerank_weights_by_position():
    # 12 candidates, retrieval scores 12 down to 1 → norms 1.0 down to 0.0.
    candidates = [(cid, float(12 - i)) for i, cid in enumerate(range(100, 112))]
    details = {cid: {} for cid, _ in candidates}
    rr = [1.0] * len(candidates)

    blended = dict(_blend_rerank(candidates, rr, details))

    assert details[100]["blend_weight"] == 0.75  # rank 1
    assert details[102]["blend_weight"] == 0.75  # rank 3
    assert details[103]["blend_weight"] == 0.60  # rank 4
    assert details[109]["blend_weight"] == 0.60  # rank 10
    assert details[110]["blend_weight"] == 0.40  # rank 11

    # rank 1: norm 1.0 → 0.75*1.0 + 0.25*1.0 = 1.0
    assert blended[100] == pytest.approx(1.0)
    # rank 4: norm 8/11 → 0.6*(8/11) + 0.4*1.0
    assert blended[103] == pytest.approx(0.6 * (8 / 11) + 0.4)
    # rank 11: norm 1/11 → 0.4*(1/11) + 0.6*1.0
    assert blended[110] == pytest.approx(0.4 * (1 / 11) + 0.6)


def test_blend_rerank_constant_retrieval_scores():
    """Zero span (all candidates same retrieval score) must not divide by zero."""
    candidates = [(1, 0.5), (2, 0.5)]
    details = {1: {}, 2: {}}
    blended = dict(_blend_rerank(candidates, [0.2, 0.8], details))
    assert blended[1] == pytest.approx(0.75 * 1.0 + 0.25 * 0.2)
    assert blended[2] == pytest.approx(0.75 * 1.0 + 0.25 * 0.8)