"""
RBSP Result Diversification.

Implements Maximal Marginal Relevance (MMR) reranking and category-based
diversification to prevent result lists dominated by a single file,
directory, or topic cluster.

MMR formula:
    score(d) = lambda * relevance(d) - (1 - lambda) * max_sim(d, selected)

Higher lambda favours relevance; lower lambda favours diversity.
"""

from __future__ import annotations

import heapq
from typing import Any, Callable, Dict, List, Optional, Set, Tuple


def mmr_rerank(
    results: list[tuple[str, float]],
    similarity_fn: Callable[[str, str], float],
    lambda_param: float = 0.5,
    top_k: int = 10,
) -> list[tuple[str, float]]:
    """Rerank results using Maximal Marginal Relevance.

    Greedily selects documents that maximise the trade-off between
    relevance (original score) and diversity (dissimilarity to
    already-selected documents).

    Args:
        results:       List of (item_id, relevance_score) pairs.
        similarity_fn: Callable(item_a, item_b) -> float in [0, 1].
        lambda_param:  Balance factor in [0, 1].  1.0 = pure relevance,
                       0.0 = pure diversity.
        top_k:         Maximum number of results to return.

    Returns:
        Reranked list of (item_id, mmr_score) pairs.
    """
    if not results:
        return []

    top_k = min(top_k, len(results))

    selected: list[tuple[str, float]] = []
    selected_ids: list[str] = []
    remaining: dict[str, float] = {item_id: score for item_id, score in results}

    for _ in range(top_k):
        best_id: Optional[str] = None
        best_mmr = float("-inf")

        for item_id, relevance in remaining.items():
            # Max similarity to any already-selected document
            if selected_ids:
                max_sim = max(
                    similarity_fn(item_id, sel_id) for sel_id in selected_ids
                )
            else:
                max_sim = 0.0

            mmr_score = lambda_param * relevance - (1.0 - lambda_param) * max_sim

            if mmr_score > best_mmr:
                best_mmr = mmr_score
                best_id = item_id

        if best_id is None:
            break

        selected.append((best_id, best_mmr))
        selected_ids.append(best_id)
        del remaining[best_id]

    return selected


class DiversityReranker:
    """Reranks search results for diversity using MMR with configurable
    similarity metrics.

    Supports path-based similarity (files in the same directory are
    considered similar), token-set Jaccard similarity, and a weighted
    combination of both.

    Args:
        lambda_param:     MMR balance factor in [0, 1].
        similarity_metric: One of ``"path"``, ``"token"``, or ``"combined"``.
        path_weight:      Weight of path similarity in combined mode (default 0.6).
    """

    def __init__(
        self,
        lambda_param: float = 0.5,
        similarity_metric: str = "path",
        path_weight: float = 0.6,
    ) -> None:
        self.lambda_param = lambda_param
        self.similarity_metric = similarity_metric
        self.path_weight = path_weight

        # Caches populated during rerank for similarity lookups
        self._token_cache: Dict[str, set[str]] = {}

    # ------------------------------------------------------------------
    # Similarity functions
    # ------------------------------------------------------------------

    @staticmethod
    def path_similarity(path_a: str, path_b: str) -> float:
        """Similarity based on shared path components.

        Files in the same directory receive high similarity; files sharing
        no path components receive 0.0.

        Returns:
            Float in [0.0, 1.0].
        """
        parts_a = _split_path(path_a)
        parts_b = _split_path(path_b)

        if not parts_a or not parts_b:
            return 0.0

        # Count leading shared components
        shared = 0
        for a, b in zip(parts_a, parts_b):
            if a == b:
                shared += 1
            else:
                break

        max_depth = max(len(parts_a), len(parts_b))
        return shared / max_depth if max_depth > 0 else 0.0

    @staticmethod
    def token_similarity(tokens_a: set[str], tokens_b: set[str]) -> float:
        """Jaccard similarity of two token sets.

        Returns:
            Float in [0.0, 1.0].
        """
        if not tokens_a and not tokens_b:
            return 1.0
        if not tokens_a or not tokens_b:
            return 0.0
        intersection = len(tokens_a & tokens_b)
        union = len(tokens_a | tokens_b)
        return intersection / union if union > 0 else 0.0

    def combined_similarity(self, a: str, b: str) -> float:
        """Weighted combination of path and token similarity.

        Uses ``path_weight`` for path component and ``1 - path_weight``
        for token-set Jaccard.

        Requires that token sets have been populated in ``_token_cache``
        (done automatically by :meth:`rerank`).
        """
        p_sim = self.path_similarity(a, b)
        tokens_a = self._token_cache.get(a, set())
        tokens_b = self._token_cache.get(b, set())
        t_sim = self.token_similarity(tokens_a, tokens_b)
        return self.path_weight * p_sim + (1.0 - self.path_weight) * t_sim

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

    def rerank(
        self,
        results: list[Any],
        top_k: int = 10,
        id_fn: Optional[Callable[[Any], str]] = None,
        score_fn: Optional[Callable[[Any], float]] = None,
        tokens_fn: Optional[Callable[[Any], set[str]]] = None,
    ) -> list[Any]:
        """Rerank *results* for diversity using MMR.

        The caller may pass accessor functions so that this method works
        with arbitrary result objects.  By default results are assumed to
        be ``(item_id, score)`` tuples.

        Args:
            results:   The initial ranked result list.
            top_k:     Maximum results to return.
            id_fn:     Extracts a string identifier from a result.
            score_fn:  Extracts the relevance score from a result.
            tokens_fn: Extracts a set of string tokens from a result
                       (needed for ``"token"`` and ``"combined"`` metrics).

        Returns:
            Reranked list of result objects (same type as input).
        """
        if not results:
            return []

        # Default accessors for (id, score) tuples
        if id_fn is None:
            id_fn = lambda r: r[0]
        if score_fn is None:
            score_fn = lambda r: r[1]

        # Build lookup structures
        result_map: Dict[str, Any] = {}
        relevance_map: Dict[str, float] = {}
        self._token_cache.clear()

        for r in results:
            rid = id_fn(r)
            result_map[rid] = r
            relevance_map[rid] = score_fn(r)
            if tokens_fn is not None:
                self._token_cache[rid] = tokens_fn(r)

        # Select the right similarity function
        sim_fn = self._get_similarity_fn()

        # Use the core MMR algorithm
        pairs = [(rid, rel) for rid, rel in relevance_map.items()]
        reranked = mmr_rerank(
            pairs,
            similarity_fn=sim_fn,
            lambda_param=self.lambda_param,
            top_k=top_k,
        )

        return [result_map[rid] for rid, _score in reranked]

    def _get_similarity_fn(self) -> Callable[[str, str], float]:
        if self.similarity_metric == "token":
            return lambda a, b: self.token_similarity(
                self._token_cache.get(a, set()),
                self._token_cache.get(b, set()),
            )
        if self.similarity_metric == "combined":
            return self.combined_similarity
        # Default: path
        return self.path_similarity


class CategoryDiversifier:
    """Ensures results span multiple categories.

    Given a category extraction function, guarantees that at least
    ``min_per_category`` results from each represented category appear
    in the output (when possible), then fills remaining slots by
    relevance.

    This is complementary to MMR: use MMR for continuous similarity
    diversification and ``CategoryDiversifier`` for hard categorical
    guarantees.
    """

    def diversify(
        self,
        results: list[Any],
        category_fn: Callable[[Any], str],
        min_per_category: int = 1,
        top_k: int = 10,
        score_fn: Optional[Callable[[Any], float]] = None,
    ) -> list[Any]:
        """Diversify *results* across categories.

        Args:
            results:          The initial ranked result list.
            category_fn:      Callable that returns a category string for
                              a result object.
            min_per_category: Minimum results to include per category.
            top_k:            Maximum results to return.
            score_fn:         Extracts relevance score (for tie-breaking).
                              Defaults to ``lambda r: r[1]``.

        Returns:
            Diversified list of result objects.
        """
        if not results:
            return []

        if score_fn is None:
            score_fn = lambda r: r[1]

        # Group by category, preserving original order within each group
        categories: Dict[str, list[Any]] = {}
        for r in results:
            cat = category_fn(r)
            if cat not in categories:
                categories[cat] = []
            categories[cat].append(r)

        selected: list[Any] = []
        selected_set: set[int] = set()  # track by id(r)

        # Phase 1: guarantee min_per_category from each category
        for cat, cat_results in categories.items():
            count = 0
            for r in cat_results:
                if count >= min_per_category:
                    break
                if id(r) not in selected_set:
                    selected.append(r)
                    selected_set.add(id(r))
                    count += 1

        # Phase 2: fill remaining slots by relevance
        if len(selected) < top_k:
            remaining = [r for r in results if id(r) not in selected_set]
            # Sort remaining by score descending
            remaining.sort(key=lambda r: score_fn(r), reverse=True)
            for r in remaining:
                if len(selected) >= top_k:
                    break
                selected.append(r)
                selected_set.add(id(r))

        return selected[:top_k]


# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------

def _split_path(path: str) -> list[str]:
    """Split a path into components, handling both / and \\ separators."""
    # Normalise to forward slashes
    normalized = path.replace("\\", "/")
    parts = [p for p in normalized.split("/") if p]
    return parts
