"""Tests for RBSP result diversification (MMR + category diversifier)."""

from __future__ import annotations

import pytest

from rbsp.query.diversity import (
    CategoryDiversifier,
    DiversityReranker,
    mmr_rerank,
    _split_path,
)


# ── helpers ──────────────────────────────────────────────────────────

def _constant_sim(a: str, b: str) -> float:
    """All items are equally similar."""
    return 0.5 if a != b else 1.0


def _zero_sim(a: str, b: str) -> float:
    """No items are similar."""
    return 0.0


def _identity_sim(a: str, b: str) -> float:
    """Only identical items are similar."""
    return 1.0 if a == b else 0.0


# ── mmr_rerank ───────────────────────────────────────────────────────

class TestMmrRerank:
    def test_empty_input(self):
        assert mmr_rerank([], _zero_sim) == []

    def test_single_result(self):
        results = [("a", 0.9)]
        out = mmr_rerank(results, _zero_sim, top_k=5)
        assert len(out) == 1
        assert out[0][0] == "a"

    def test_top_k_limits_output(self):
        results = [("a", 0.9), ("b", 0.8), ("c", 0.7), ("d", 0.6)]
        out = mmr_rerank(results, _zero_sim, top_k=2)
        assert len(out) == 2

    def test_pure_relevance_lambda_1(self):
        """lambda=1.0 should rank purely by relevance."""
        results = [("a", 0.5), ("b", 0.9), ("c", 0.7)]
        out = mmr_rerank(results, _constant_sim, lambda_param=1.0, top_k=3)
        ids = [item_id for item_id, _ in out]
        assert ids == ["b", "c", "a"]

    def test_diversity_penalises_similar_items(self):
        """With low lambda, similar items should be penalised."""
        # Create results where a and b are very similar but c is different
        def selective_sim(a: str, b: str) -> float:
            if {a, b} == {"a", "b"}:
                return 0.95  # a and b are very similar
            return 0.0

        results = [("a", 0.9), ("b", 0.89), ("c", 0.5)]
        out = mmr_rerank(results, selective_sim, lambda_param=0.3, top_k=3)
        ids = [item_id for item_id, _ in out]
        # a should be first (highest relevance), then c (dissimilar to a),
        # then b (very similar to a, heavily penalised)
        assert ids[0] == "a"
        assert ids[1] == "c"
        assert ids[2] == "b"

    def test_zero_similarity_preserves_relevance_order(self):
        """When all items are dissimilar, relevance order is preserved."""
        results = [("a", 0.9), ("b", 0.7), ("c", 0.5)]
        out = mmr_rerank(results, _zero_sim, lambda_param=0.5, top_k=3)
        ids = [item_id for item_id, _ in out]
        assert ids == ["a", "b", "c"]

    def test_top_k_exceeds_results(self):
        results = [("a", 0.9), ("b", 0.8)]
        out = mmr_rerank(results, _zero_sim, top_k=100)
        assert len(out) == 2


# ── _split_path ──────────────────────────────────────────────────────

class TestSplitPath:
    def test_unix_path(self):
        assert _split_path("src/rbsp/query/ranker.py") == [
            "src", "rbsp", "query", "ranker.py"
        ]

    def test_windows_path(self):
        assert _split_path("src\\rbsp\\query\\ranker.py") == [
            "src", "rbsp", "query", "ranker.py"
        ]

    def test_empty(self):
        assert _split_path("") == []

    def test_root_slash(self):
        assert _split_path("/usr/bin/python") == ["usr", "bin", "python"]


# ── DiversityReranker.path_similarity ────────────────────────────────

class TestPathSimilarity:
    def test_same_file(self):
        sim = DiversityReranker.path_similarity(
            "src/rbsp/query/ranker.py",
            "src/rbsp/query/ranker.py",
        )
        assert sim == 1.0

    def test_same_directory(self):
        sim = DiversityReranker.path_similarity(
            "src/rbsp/query/ranker.py",
            "src/rbsp/query/diversity.py",
        )
        # Share 3 of 4 leading components: src, rbsp, query
        assert 0.7 <= sim <= 0.8

    def test_different_top_directory(self):
        sim = DiversityReranker.path_similarity(
            "src/rbsp/query/ranker.py",
            "tests/test_ranker.py",
        )
        assert sim == 0.0

    def test_empty_path(self):
        assert DiversityReranker.path_similarity("", "src/foo.py") == 0.0
        assert DiversityReranker.path_similarity("src/foo.py", "") == 0.0

    def test_partial_overlap(self):
        sim = DiversityReranker.path_similarity(
            "src/rbsp/core/lsh.py",
            "src/rbsp/query/ranker.py",
        )
        # Share 2 of 4: src, rbsp
        assert 0.4 <= sim <= 0.6


# ── DiversityReranker.token_similarity ───────────────────────────────

class TestTokenSimilarity:
    def test_identical_sets(self):
        s = {"hello", "world"}
        assert DiversityReranker.token_similarity(s, s) == 1.0

    def test_disjoint_sets(self):
        assert DiversityReranker.token_similarity({"a"}, {"b"}) == 0.0

    def test_partial_overlap(self):
        sim = DiversityReranker.token_similarity(
            {"a", "b", "c"}, {"b", "c", "d"}
        )
        # Jaccard: 2/4 = 0.5
        assert sim == 0.5

    def test_both_empty(self):
        assert DiversityReranker.token_similarity(set(), set()) == 1.0

    def test_one_empty(self):
        assert DiversityReranker.token_similarity({"a"}, set()) == 0.0


# ── DiversityReranker.rerank ─────────────────────────────────────────

class TestDiversityRerankerRerank:
    def test_empty(self):
        r = DiversityReranker()
        assert r.rerank([]) == []

    def test_basic_path_diversity(self):
        """Results from same directory should be spread out."""
        results = [
            ("src/query/a.py", 0.9),
            ("src/query/b.py", 0.85),
            ("src/query/c.py", 0.80),
            ("src/core/d.py", 0.6),
            ("tests/e.py", 0.5),
        ]
        r = DiversityReranker(lambda_param=0.4, similarity_metric="path")
        out = r.rerank(results, top_k=5)
        ids = [item[0] for item in out]
        # First result should be highest relevance
        assert ids[0] == "src/query/a.py"
        # With low lambda, d.py and e.py (different dirs) should be boosted
        # relative to b.py and c.py (same dir as a.py)
        # At minimum, we shouldn't have all query/ files clustered at the top
        non_query_positions = [
            i for i, item_id in enumerate(ids)
            if not item_id.startswith("src/query/")
        ]
        assert len(non_query_positions) > 0

    def test_custom_accessors(self):
        """rerank works with custom id_fn and score_fn."""
        results = [
            {"path": "a.py", "score": 0.9},
            {"path": "b.py", "score": 0.8},
        ]
        r = DiversityReranker(similarity_metric="path")
        out = r.rerank(
            results,
            top_k=2,
            id_fn=lambda x: x["path"],
            score_fn=lambda x: x["score"],
        )
        assert len(out) == 2
        assert all(isinstance(item, dict) for item in out)

    def test_token_similarity_metric(self):
        results = [
            ("a.py", 0.9),
            ("b.py", 0.85),
            ("c.py", 0.7),
        ]
        r = DiversityReranker(
            lambda_param=0.5,
            similarity_metric="token",
        )
        out = r.rerank(
            results,
            top_k=3,
            tokens_fn=lambda item: {"hello", "world"} if item[0] in ("a.py", "b.py") else {"foo"},
        )
        ids = [item[0] for item in out]
        assert ids[0] == "a.py"
        # c.py should be boosted because its tokens differ from a.py
        assert ids[1] == "c.py"

    def test_combined_similarity_metric(self):
        results = [
            ("src/a.py", 0.9),
            ("src/b.py", 0.85),
            ("lib/c.py", 0.7),
        ]
        r = DiversityReranker(
            lambda_param=0.5,
            similarity_metric="combined",
            path_weight=0.5,
        )
        out = r.rerank(
            results,
            top_k=3,
            tokens_fn=lambda item: {"x", "y"} if item[0].startswith("src") else {"z"},
        )
        assert len(out) == 3
        ids = [item[0] for item in out]
        assert ids[0] == "src/a.py"


# ── CategoryDiversifier ─────────────────────────────────────────────

class TestCategoryDiversifier:
    def _dir_category(self, item):
        """Extract directory as category from (path, score) tuples."""
        path = item[0]
        parts = path.replace("\\", "/").rsplit("/", 1)
        return parts[0] if len(parts) > 1 else ""

    def test_empty(self):
        d = CategoryDiversifier()
        assert d.diversify([], category_fn=lambda x: "a") == []

    def test_single_category(self):
        results = [
            ("src/a.py", 0.9),
            ("src/b.py", 0.8),
            ("src/c.py", 0.7),
        ]
        d = CategoryDiversifier()
        out = d.diversify(results, category_fn=self._dir_category, top_k=2)
        assert len(out) == 2

    def test_guarantees_min_per_category(self):
        results = [
            ("src/a.py", 0.9),
            ("src/b.py", 0.85),
            ("src/c.py", 0.80),
            ("tests/d.py", 0.3),
            ("lib/e.py", 0.2),
        ]
        d = CategoryDiversifier()
        out = d.diversify(
            results,
            category_fn=self._dir_category,
            min_per_category=1,
            top_k=5,
        )
        cats_in_output = {self._dir_category(item) for item in out}
        # All three categories should appear
        assert "src" in cats_in_output
        assert "tests" in cats_in_output
        assert "lib" in cats_in_output

    def test_min_per_category_respected(self):
        results = [
            ("src/a.py", 0.9),
            ("tests/b.py", 0.85),
            ("src/c.py", 0.8),
            ("tests/d.py", 0.75),
        ]
        d = CategoryDiversifier()
        out = d.diversify(
            results,
            category_fn=self._dir_category,
            min_per_category=2,
            top_k=4,
        )
        from collections import Counter
        cat_counts = Counter(self._dir_category(item) for item in out)
        assert cat_counts["src"] >= 2
        assert cat_counts["tests"] >= 2

    def test_top_k_limits(self):
        results = [(f"dir{i}/f.py", 1.0 - i * 0.1) for i in range(20)]
        d = CategoryDiversifier()
        out = d.diversify(
            results,
            category_fn=self._dir_category,
            top_k=5,
        )
        assert len(out) == 5

    def test_fills_remaining_by_relevance(self):
        results = [
            ("src/a.py", 0.9),
            ("tests/b.py", 0.5),
            ("src/c.py", 0.85),
            ("src/d.py", 0.80),
        ]
        d = CategoryDiversifier()
        out = d.diversify(
            results,
            category_fn=self._dir_category,
            min_per_category=1,
            top_k=4,
        )
        # Phase 1: src/a.py (1st from src), tests/b.py (1st from tests)
        # Phase 2: fill remaining 2 slots by relevance: src/c.py (0.85), src/d.py (0.80)
        assert len(out) == 4
        ids = [item[0] for item in out]
        assert "src/a.py" in ids
        assert "tests/b.py" in ids


# ── DiversityReranker.combined_similarity ────────────────────────────

class TestCombinedSimilarity:
    def test_weights_applied(self):
        r = DiversityReranker(path_weight=0.5)
        r._token_cache = {
            "src/a.py": {"hello", "world"},
            "src/b.py": {"hello", "world"},
        }
        sim = r.combined_similarity("src/a.py", "src/b.py")
        p_sim = DiversityReranker.path_similarity("src/a.py", "src/b.py")
        t_sim = DiversityReranker.token_similarity(
            {"hello", "world"}, {"hello", "world"}
        )
        expected = 0.5 * p_sim + 0.5 * t_sim
        assert abs(sim - expected) < 1e-9

    def test_missing_tokens_treated_as_empty(self):
        r = DiversityReranker(path_weight=0.5)
        r._token_cache = {}
        sim = r.combined_similarity("src/a.py", "lib/b.py")
        # token sim of empty sets = 1.0, path sim = 0.0
        expected = 0.5 * 0.0 + 0.5 * 1.0
        assert abs(sim - expected) < 1e-9
