"""Tests for search module RRF fusion functions.

Tests are written against the contracts from src/lattice/core/search.py.
These tests MUST FAIL initially because implementations are stubs.
"""

import pytest

from lattice.core.search import (
    RRF_K,
    assign_positional_ranks,
    calculate_rrf_score,
    fuse_results,
    get_unique_rowids,
    top_n,
)
from lattice.core.types.search import FtsResult, RankedItem, VecResult


class TestRRFConstant:
    """Tests for RRF_K constant."""

    def test_rrf_k_value(self) -> None:
        """RRF_K should be 60 per RFC-002."""
        assert RRF_K == 60, f"Expected RRF_K=60, got {RRF_K}"


class TestAssignPositionalRanks:
    """Tests for assign_positional_ranks function."""

    def test_returns_dict(self) -> None:
        """Should return a dict (contract: isinstance(result, dict))."""
        fts_results = [FtsResult(rowid=1, rank=-0.5)]
        result = assign_positional_ranks(fts_results)
        assert isinstance(result, dict), f"Expected dict, got {type(result)}"

    def test_keys_are_integers(self) -> None:
        """All keys should be integers (contract: all isinstance(k, int))."""
        fts_results = [
            FtsResult(rowid=1, rank=-0.5),
            FtsResult(rowid=2, rank=-0.3),
        ]
        result = assign_positional_ranks(fts_results)
        for k in result.keys():
            assert isinstance(k, int), f"Expected int key, got {type(k)}: {k}"

    def test_values_are_integers(self) -> None:
        """All values should be integers (contract: all isinstance(v, int))."""
        fts_results = [
            FtsResult(rowid=1, rank=-0.5),
            FtsResult(rowid=2, rank=-0.3),
        ]
        result = assign_positional_ranks(fts_results)
        for v in result.values():
            assert isinstance(v, int), f"Expected int value, got {type(v)}: {v}"

    def test_fts_results_positional_rank_ordering(self) -> None:
        """For FTS: lower rank (more negative) = higher position (rank 1)."""
        # More negative rank means better match in BM25
        fts_results = [
            FtsResult(rowid=1, rank=-0.9),  # Best match -> position 1
            FtsResult(rowid=2, rank=-0.5),  # Second best -> position 2
            FtsResult(rowid=3, rank=-0.1),  # Worst -> position 3
        ]
        result = assign_positional_ranks(fts_results)
        assert result[1] == 1, f"Best FTS match should have rank 1, got {result[1]}"
        assert result[2] == 2, f"Second FTS match should have rank 2, got {result[2]}"
        assert result[3] == 3, f"Third FTS match should have rank 3, got {result[3]}"

    def test_vec_results_positional_rank_ordering(self) -> None:
        """For Vec: lower distance = higher position (rank 1)."""
        vec_results = [
            VecResult(rowid=1, distance=0.1),  # Best match -> position 1
            VecResult(rowid=2, distance=0.5),  # Second best -> position 2
            VecResult(rowid=3, distance=0.9),  # Worst -> position 3
        ]
        result = assign_positional_ranks(vec_results)
        assert result[1] == 1, f"Best Vec match should have rank 1, got {result[1]}"
        assert result[2] == 2, f"Second Vec match should have rank 2, got {result[2]}"
        assert result[3] == 3, f"Third Vec match should have rank 3, got {result[3]}"

    def test_empty_fts_results(self) -> None:
        """Empty FTS list should return empty dict."""
        result = assign_positional_ranks([])
        assert result == {}, f"Expected empty dict, got {result}"

    def test_empty_vec_results(self) -> None:
        """Empty Vec list should return empty dict."""
        result = assign_positional_ranks([])  # type: ignore[arg-type]
        assert result == {}, f"Expected empty dict, got {result}"

    def test_ranks_are_one_based(self) -> None:
        """Positional ranks should be 1-based, not 0-based."""
        fts_results = [FtsResult(rowid=1, rank=-0.5)]
        result = assign_positional_ranks(fts_results)
        assert result[1] >= 1, f"Ranks should be 1-based, got {result[1]}"

    def test_single_result_gets_rank_one(self) -> None:
        """Single result should get rank 1."""
        fts_results = [FtsResult(rowid=42, rank=-0.5)]
        result = assign_positional_ranks(fts_results)
        assert result[42] == 1, f"Single result should have rank 1, got {result[42]}"


class TestCalculateRRFScore:
    """Tests for calculate_rrf_score function."""

    def test_returns_non_negative(self) -> None:
        """Should return non-negative result (contract: result >= 0)."""
        result = calculate_rrf_score(bm25_rank=1, vec_rank=1)
        assert result >= 0, f"Expected non-negative result, got {result}"

    def test_both_ranks_present(self) -> None:
        """When both ranks are present, score = 1/(k+r1) + 1/(k+r2)."""
        # RRF Score = 1/(k + bm25_rank) + 1/(k + vec_rank)
        # k=60, bm25_rank=1, vec_rank=2
        # Expected: 1/(60+1) + 1/(60+2) = 1/61 + 1/62
        result = calculate_rrf_score(bm25_rank=1, vec_rank=2)
        expected = 1 / 61 + 1 / 62
        assert abs(result - expected) < 0.0001, f"Expected {expected}, got {result}"

    def test_bm25_rank_only(self) -> None:
        """When only BM25 rank is present (vec_rank=None)."""
        # Should still work with partial data
        result = calculate_rrf_score(bm25_rank=1, vec_rank=None)
        # With vec_rank=None, the contribution from vector should be 0
        # Expected: 1/(k+1) + 0 = 1/61
        expected_bm25_only = 1 / 61
        # The contract doesn't specify exact behavior for None,
        # but it should return a valid non-negative score
        assert result >= 0, f"Expected non-negative score, got {result}"

    def test_vec_rank_only(self) -> None:
        """When only vector rank is present (bm25_rank=None)."""
        result = calculate_rrf_score(bm25_rank=None, vec_rank=1)
        # Should return a valid non-negative score
        assert result >= 0, f"Expected non-negative score, got {result}"

    def test_neither_rank_present(self) -> None:
        """When both ranks are None."""
        result = calculate_rrf_score(bm25_rank=None, vec_rank=None)
        # Should return 0 when no ranks are available
        assert result >= 0, f"Expected non-negative score, got {result}"

    def test_custom_k_parameter(self) -> None:
        """Should accept custom k parameter."""
        # With k=100, bm25_rank=1, vec_rank=2
        # Expected: 1/(100+1) + 1/(100+2)
        result = calculate_rrf_score(bm25_rank=1, vec_rank=2, k=100)
        expected = 1 / 101 + 1 / 102
        assert abs(result - expected) < 0.0001, f"Expected {expected}, got {result}"

    def test_k_must_be_positive(self) -> None:
        """k parameter must be > 0 (contract: k > 0)."""
        with pytest.raises(Exception):  # deal.pre should raise
            calculate_rrf_score(bm25_rank=1, vec_rank=1, k=0)

    def test_result_is_float(self) -> None:
        """Result should be a float."""
        result = calculate_rrf_score(bm25_rank=1, vec_rank=1)
        assert isinstance(result, float), f"Expected float, got {type(result)}"

    def test_both_ranks_higher_score_than_single(self) -> None:
        """Result appearing in both lists should have higher score than single list."""
        both_score = calculate_rrf_score(bm25_rank=1, vec_rank=1)
        bm25_only_score = calculate_rrf_score(bm25_rank=1, vec_rank=None)
        vec_only_score = calculate_rrf_score(bm25_rank=None, vec_rank=1)

        # Both present gives full RRF score
        # Single present gives partial score
        # This test verifies the behavior (exact formula may vary)
        assert both_score >= bm25_only_score or both_score >= vec_only_score, (
            f"Both ranks ({both_score}) should score >= single rank "
            f"(bm25: {bm25_only_score}, vec: {vec_only_score})"
        )


class TestFuseResults:
    """Tests for fuse_results function."""

    def test_returns_list(self) -> None:
        """Should return a list (contract: isinstance(result, list))."""
        bm25_results = [FtsResult(rowid=1, rank=-0.5)]
        vec_results = [VecResult(rowid=1, distance=0.1)]
        result = fuse_results(bm25_results, vec_results)
        assert isinstance(result, list), f"Expected list, got {type(result)}"

    def test_all_items_are_ranked_item(self) -> None:
        """All items should be RankedItem instances."""
        bm25_results = [FtsResult(rowid=1, rank=-0.5)]
        vec_results = [VecResult(rowid=2, distance=0.1)]
        result = fuse_results(bm25_results, vec_results)
        for item in result:
            assert isinstance(item, RankedItem), (
                f"Expected RankedItem, got {type(item)}"
            )

    def test_no_duplicate_rowids(self) -> None:
        """Each rowid should appear only once (contract: len(result) == len(set(rowids)))."""
        # Same rowid in both lists should be fused, not duplicated
        bm25_results = [FtsResult(rowid=1, rank=-0.5)]
        vec_results = [VecResult(rowid=1, distance=0.1)]
        result = fuse_results(bm25_results, vec_results)
        rowids = [item.rowid for item in result]
        assert len(rowids) == len(set(rowids)), f"Duplicate rowids found: {rowids}"

    def test_empty_input_lists(self) -> None:
        """Empty both input lists should return empty list."""
        result = fuse_results([], [])
        assert result == [], f"Expected empty list, got {result}"

    def test_empty_bm25_only_vec_results(self) -> None:
        """Empty BM25 with vector results should return ranked vector results."""
        vec_results = [
            VecResult(rowid=1, distance=0.1),
            VecResult(rowid=2, distance=0.5),
        ]
        result = fuse_results([], vec_results)
        assert len(result) == 2, f"Expected 2 results, got {len(result)}"
        assert all(isinstance(item, RankedItem) for item in result)

    def test_empty_vec_only_bm25_results(self) -> None:
        """Empty vector with BM25 results should return ranked BM25 results."""
        bm25_results = [
            FtsResult(rowid=1, rank=-0.5),
            FtsResult(rowid=2, rank=-0.3),
        ]
        result = fuse_results(bm25_results, [])
        assert len(result) == 2, f"Expected 2 results, got {len(result)}"
        assert all(isinstance(item, RankedItem) for item in result)

    def test_fuses_common_rowids(self) -> None:
        """Same rowid in both lists should appear once with combined score."""
        bm25_results = [FtsResult(rowid=42, rank=-0.5)]
        vec_results = [VecResult(rowid=42, distance=0.1)]
        result = fuse_results(bm25_results, vec_results)
        # Should have exactly one result with rowid 42
        assert len(result) == 1, f"Expected 1 fused result, got {len(result)}"
        assert result[0].rowid == 42, f"Expected rowid 42, got {result[0].rowid}"

    def test_sorted_by_rrf_score_descending(self) -> None:
        """Results should be sorted by RRF score in descending order."""
        # Create results where rowid 1 appears in both lists (higher combined score)
        # and rowid 2, 3 appear in only one list
        bm25_results = [
            FtsResult(rowid=1, rank=-0.9),  # Best BM25 -> rank 1
            FtsResult(rowid=2, rank=-0.5),  # rank 2
        ]
        vec_results = [
            VecResult(rowid=1, distance=0.1),  # Best Vec -> rank 1
            VecResult(rowid=3, distance=0.5),  # rank 2
        ]
        result = fuse_results(bm25_results, vec_results)
        # rowid 1 should be first (appears in both, highest combined score)
        # The list should be sorted by rrf_score descending
        scores = [item.rrf_score for item in result]
        assert scores == sorted(scores, reverse=True), (
            f"Results not sorted by RRF score descending: {scores}"
        )

    def test_custom_k_parameter(self) -> None:
        """Should accept custom k parameter."""
        bm25_results = [FtsResult(rowid=1, rank=-0.5)]
        vec_results = [VecResult(rowid=1, distance=0.1)]
        result = fuse_results(bm25_results, vec_results, k=100)
        assert isinstance(result, list), "Should accept custom k parameter"


class TestTopN:
    """Tests for top_n function."""

    def test_returns_list(self) -> None:
        """Should return a list (contract: isinstance(result, list))."""
        items = [
            RankedItem(rowid=1, rrf_score=0.05),
            RankedItem(rowid=2, rrf_score=0.03),
        ]
        result = top_n(items, 1)
        assert isinstance(result, list), f"Expected list, got {type(result)}"

    def test_all_items_are_ranked_item(self) -> None:
        """All items should still be RankedItem instances."""
        items = [
            RankedItem(rowid=1, rrf_score=0.05),
            RankedItem(rowid=2, rrf_score=0.03),
        ]
        result = top_n(items, 2)
        for item in result:
            assert isinstance(item, RankedItem), (
                f"Expected RankedItem, got {type(item)}"
            )

    def test_top_one_returns_one_item(self) -> None:
        """top_n(..., 1) should return exactly one item from sorted list."""
        items = [
            RankedItem(rowid=1, rrf_score=0.05),  # Highest
            RankedItem(rowid=2, rrf_score=0.03),
            RankedItem(rowid=3, rrf_score=0.01),
        ]
        result = top_n(items, 1)
        assert len(result) == 1, f"Expected 1 item, got {len(result)}"
        assert result[0].rowid == 1, (
            f"Expected rowid 1 (highest score), got {result[0].rowid}"
        )
        assert result[0].rrf_score == 0.05, (
            f"Expected score 0.05, got {result[0].rrf_score}"
        )

    def test_top_n_larger_than_list(self) -> None:
        """Requesting more items than available should return all items."""
        items = [
            RankedItem(rowid=1, rrf_score=0.05),
            RankedItem(rowid=2, rrf_score=0.03),
        ]
        result = top_n(items, 10)
        assert len(result) == 2, f"Expected 2 items, got {len(result)}"

    def test_top_zero_returns_empty(self) -> None:
        """top_n(..., 0) should return empty list."""
        items = [
            RankedItem(rowid=1, rrf_score=0.05),
        ]
        result = top_n(items, 0)
        assert result == [], f"Expected empty list, got {result}"

    def test_empty_input_returns_empty(self) -> None:
        """Empty input list should return empty list."""
        result = top_n([], 5)
        assert result == [], f"Expected empty list, got {result}"

    def test_preserves_sorted_order(self) -> None:
        """Should return items in the order given (input should be pre-sorted)."""
        # Input is already sorted by RRF score descending
        items = [
            RankedItem(rowid=1, rrf_score=0.05),
            RankedItem(rowid=2, rrf_score=0.03),
            RankedItem(rowid=3, rrf_score=0.01),
        ]
        result = top_n(items, 2)
        assert result[0].rowid == 1, "First item should be rowid 1"
        assert result[1].rowid == 2, "Second item should be rowid 2"


class TestGetUniqueRowids:
    """Tests for get_unique_rowids function."""

    def test_returns_list(self) -> None:
        """Should return a list (contract: isinstance(result, list))."""
        bm25_results = [FtsResult(rowid=1, rank=-0.5)]
        vec_results = [VecResult(rowid=2, distance=0.1)]
        result = get_unique_rowids(bm25_results, vec_results)
        assert isinstance(result, list), f"Expected list, got {type(result)}"

    def test_no_duplicates(self) -> None:
        """Result should have no duplicates (contract: len(result) == len(set(result)))."""
        # Same rowid in both lists
        bm25_results = [FtsResult(rowid=1, rank=-0.5)]
        vec_results = [VecResult(rowid=1, distance=0.1)]
        result = get_unique_rowids(bm25_results, vec_results)
        assert len(result) == len(set(result)), f"Expected no duplicates, got {result}"

    def test_empty_both_lists(self) -> None:
        """Empty both lists should return empty list."""
        result = get_unique_rowids([], [])
        assert result == [], f"Expected empty list, got {result}"

    def test_empty_bm25_only_vec(self) -> None:
        """Empty BM25 with vector results should return vector rowids."""
        vec_results = [
            VecResult(rowid=1, distance=0.1),
            VecResult(rowid=2, distance=0.5),
        ]
        result = get_unique_rowids([], vec_results)
        assert len(result) == 2, f"Expected 2 rowids, got {len(result)}"
        assert 1 in result and 2 in result

    def test_empty_vec_only_bm25(self) -> None:
        """Empty vector with BM25 results should return BM25 rowids."""
        bm25_results = [
            FtsResult(rowid=1, rank=-0.5),
            FtsResult(rowid=2, rank=-0.3),
        ]
        result = get_unique_rowids(bm25_results, [])
        assert len(result) == 2, f"Expected 2 rowids, got {len(result)}"
        assert 1 in result and 2 in result

    def test_order_preserved_bm25_first(self) -> None:
        """Order should be BM25 rowids first, then vector rowids (per docstring)."""
        bm25_results = [
            FtsResult(rowid=1, rank=-0.5),
            FtsResult(rowid=2, rank=-0.3),
        ]
        vec_results = [
            VecResult(rowid=3, distance=0.1),
            VecResult(rowid=4, distance=0.5),
        ]
        result = get_unique_rowids(bm25_results, vec_results)
        # BM25 rowids should come before vector rowids
        assert result.index(1) < result.index(3), (
            "BM25 rowid 1 should come before vec rowid 3"
        )
        assert result.index(1) < result.index(4), (
            "BM25 rowid 1 should come before vec rowid 4"
        )
        assert result.index(2) < result.index(3), (
            "BM25 rowid 2 should come before vec rowid 3"
        )

    def test_duplicate_rowid_appears_once(self) -> None:
        """Duplicate rowid in both lists should appear only once."""
        bm25_results = [
            FtsResult(rowid=1, rank=-0.5),
            FtsResult(rowid=2, rank=-0.3),
        ]
        vec_results = [
            VecResult(rowid=1, distance=0.1),  # Duplicate with BM25
            VecResult(rowid=3, distance=0.5),
        ]
        result = get_unique_rowids(bm25_results, vec_results)
        assert result.count(1) == 1, (
            f"Rowid 1 should appear exactly once, got {result.count(1)}"
        )
        assert len(result) == 3, (
            f"Expected 3 unique rowids (1, 2, 3), got {len(result)}"
        )

    def test_all_rowids_are_integers(self) -> None:
        """All returned rowids should be integers."""
        bm25_results = [FtsResult(rowid=1, rank=-0.5)]
        vec_results = [VecResult(rowid=2, distance=0.1)]
        result = get_unique_rowids(bm25_results, vec_results)
        for rowid in result:
            assert isinstance(rowid, int), (
                f"Expected int rowid, got {type(rowid)}: {rowid}"
            )


class TestIntegrationScenarios:
    """Integration tests exercising multiple functions together.

    These tests verify the complete RRF fusion workflow.
    """

    def test_full_rrf_workflow(self) -> None:
        """Complete workflow: get unique rowids -> assign ranks -> calculate scores."""
        # Setup: BM25 and vector results with some overlap
        bm25_results = [
            FtsResult(rowid=1, rank=-0.9),  # Best BM25 match
            FtsResult(rowid=2, rank=-0.5),  # Second BM25 match
        ]
        vec_results = [
            VecResult(rowid=1, distance=0.1),  # Best Vec match (overlaps with BM25!)
            VecResult(rowid=3, distance=0.5),  # Second Vec match (unique)
        ]

        # Step 1: Get unique rowids
        unique_rowids = get_unique_rowids(bm25_results, vec_results)
        assert len(unique_rowids) == 3, (
            f"Expected 3 unique rowids, got {len(unique_rowids)}"
        )

        # Step 2: Fuse results using RRF
        fused = fuse_results(bm25_results, vec_results)
        assert len(fused) == 3, f"Expected 3 fused results, got {len(fused)}"

        # Step 3: Get top 2
        top_results = top_n(fused, 2)
        assert len(top_results) == 2, f"Expected 2 top results, got {len(top_results)}"

        # The rowid that appears in both lists should have highest score
        # and thus be in top_results
        top_rowids = [item.rowid for item in top_results]
        assert 1 in top_rowids, (
            f"Rowid 1 should be in top results (appears in both lists), got {top_rowids}"
        )

    def test_rrf_score_calculation_correctness(self) -> None:
        """Verify RRF score calculation matches RFC-002 formula.

        RRF Score = 1/(k + pos_bm25) + 1/(k + pos_vec), where k = 60
        """
        # A result at position 1 in both lists should have highest score
        # Score = 1/(60+1) + 1/(60+1) = 2/61 ≈ 0.0328
        score_both_position_1 = calculate_rrf_score(bm25_rank=1, vec_rank=1)
        expected_both_p1 = 1 / 61 + 1 / 61
        assert abs(score_both_position_1 - expected_both_p1) < 0.0001, (
            f"Expected {expected_both_p1}, got {score_both_position_1}"
        )

        # A result at position 1 in BM25, position 2 in vector
        # Score = 1/(60+1) + 1/(60+2) = 1/61 + 1/62 ≈ 0.0321
        score_mixed = calculate_rrf_score(bm25_rank=1, vec_rank=2)
        expected_mixed = 1 / 61 + 1 / 62
        assert abs(score_mixed - expected_mixed) < 0.0001, (
            f"Expected {expected_mixed}, got {score_mixed}"
        )

        # A result at position 2 in both lists should have lower score than position 1
        score_both_position_2 = calculate_rrf_score(bm25_rank=2, vec_rank=2)
        assert score_both_position_1 > score_both_position_2, (
            f"Position 1 ({score_both_position_1}) should score higher than position 2 ({score_both_position_2})"
        )

    def test_no_overlap_scenario(self) -> None:
        """When BM25 and vector results have no overlapping rowids."""
        bm25_results = [
            FtsResult(rowid=1, rank=-0.9),
            FtsResult(rowid=2, rank=-0.5),
        ]
        vec_results = [
            VecResult(rowid=3, distance=0.1),
            VecResult(rowid=4, distance=0.5),
        ]

        # All 4 rowids should appear
        unique_rowids = get_unique_rowids(bm25_results, vec_results)
        assert len(unique_rowids) == 4
        assert set(unique_rowids) == {1, 2, 3, 4}

        # Fuse should produce 4 results
        fused = fuse_results(bm25_results, vec_results)
        assert len(fused) == 4

    def test_all_overlap_scenario(self) -> None:
        """When all BM25 and vector results have the same rowids."""
        bm25_results = [
            FtsResult(rowid=1, rank=-0.9),
            FtsResult(rowid=2, rank=-0.5),
        ]
        vec_results = [
            VecResult(rowid=1, distance=0.1),
            VecResult(rowid=2, distance=0.5),
        ]

        # Only 2 unique rowids
        unique_rowids = get_unique_rowids(bm25_results, vec_results)
        assert len(unique_rowids) == 2
        assert set(unique_rowids) == {1, 2}

        # Fuse should produce 2 results (no duplicates)
        fused = fuse_results(bm25_results, vec_results)
        assert len(fused) == 2
        rowids = [item.rowid for item in fused]
        assert len(rowids) == len(set(rowids)), "No duplicates in fused results"


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
