//! Reciprocal Rank Fusion (RRF) for hybrid search.
//!
//! Merges ranked lists from heterogeneous retrieval sources (BM25 text search,
//! ANN vector search) into a single fused ranking. The RRF formula from
//! Cormack, Clarke & Buettcher (SIGIR 2009) is rank-based, not score-based,
//! making it robust to the incomparable score distributions of different
//! retrieval models.
//!
//! # Formula
//!
//! `RRFscore(d) = sum_i 1 / (k + rank_i(d))`
//!
//! where `k` is a smoothing constant (default 60) and `rank_i(d)` is the
//! 1-based rank of document `d` in list `i`. Documents absent from a list
//! contribute zero for that term.

use std::collections::HashMap;

use crate::{schema::EntityId, storage::vector::VectorSearchResult};

/// Reciprocal Rank Fusion (Cormack et al. SIGIR 2009).
///
/// Fuses two ranked lists into a single ranking using rank-based scoring.
/// The `k` parameter controls how much weight is given to documents ranked
/// lower in the input lists. Higher `k` compresses the score range, making
/// rank differences less significant.
#[derive(Debug, Clone)]
pub struct HybridFusion {
    /// Smoothing constant. Default is 60 per the original paper.
    pub k: u32,
}

impl Default for HybridFusion {
    fn default() -> Self {
        Self { k: 60 }
    }
}

/// Compute a single Reciprocal Rank Fusion term for a 1-based rank.
///
/// `RRFterm(rank) = 1 / (k + rank)`. Extracted so [`HybridFusion::fuse`] and
/// [`route_results`]'s `VectorOnly` path share one definition of the rank-score
/// formula rather than duplicating `1.0 / (k + rank)` at two call sites.
#[must_use]
#[inline]
pub fn rrf_term(k: f64, rank_1based: f64) -> f64 {
    1.0 / (k + rank_1based)
}

impl HybridFusion {
    /// Create a new `HybridFusion` with the default `k = 60`.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Create a new `HybridFusion` with a custom `k` value.
    #[must_use]
    pub const fn with_k(k: u32) -> Self {
        Self { k }
    }

    /// Fuse two ranked lists via Reciprocal Rank Fusion.
    ///
    /// Both lists must be pre-sorted "best first" by the caller:
    /// - `bm25_results`: sorted descending by BM25 score (index 0 = rank 1)
    /// - `ann_results`: sorted ascending by L2 distance (index 0 = rank 1)
    ///
    /// Returns results sorted descending by fused RRF score. Documents
    /// appearing in only one list contribute only their single-list term.
    #[must_use]
    #[allow(clippy::cast_precision_loss)] // Ranks are bounded by list length, never near 2^52.
    pub fn fuse(
        &self,
        bm25_results: &[(EntityId, f32)],
        ann_results: &[(EntityId, f32)],
    ) -> Vec<(EntityId, f64)> {
        let k = f64::from(self.k);
        let capacity = bm25_results.len() + ann_results.len();
        let mut scores: HashMap<u64, f64> = HashMap::with_capacity(capacity);

        for (rank_0based, (entity_id, _score)) in bm25_results.iter().enumerate() {
            let rank = (rank_0based + 1) as f64;
            *scores.entry(entity_id.as_u64()).or_insert(0.0) += rrf_term(k, rank);
        }

        for (rank_0based, (entity_id, _distance)) in ann_results.iter().enumerate() {
            let rank = (rank_0based + 1) as f64;
            *scores.entry(entity_id.as_u64()).or_insert(0.0) += rrf_term(k, rank);
        }

        let mut results: Vec<(EntityId, f64)> = scores
            .into_iter()
            .map(|(id, score)| (EntityId::new(id), score))
            .collect();
        // Sort descending by fused score, breaking ties on ASCENDING entity id so
        // the order is a true total order. The scores arrive from a `HashMap`
        // whose iteration order is process-randomized (default `RandomState`), and
        // `sort_by` is stable — without the id tie-break two entities with an equal
        // fused score would keep whatever relative order the randomized map
        // iteration produced, which changes between process runs (W14). RRF terms
        // are always finite positive, so `partial_cmp` never sees `NaN` here; the
        // only missing piece was the deterministic id tie-break.
        results.sort_by(|a, b| {
            b.1.partial_cmp(&a.1)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| a.0.as_u64().cmp(&b.0.as_u64()))
        });
        results
    }
}

/// Which retrieval system(s) to use for a search query.
///
/// Determined from the presence of text and vector components in the query.
/// Used by the retrieval router to select the appropriate fusion path.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RetrievalMode {
    /// Only text (BM25) retrieval. No embedding was provided.
    TextOnly,
    /// Only vector (ANN) retrieval. No text query was provided.
    VectorOnly,
    /// Both text and vector retrieval, fused via RRF.
    Hybrid,
}

impl RetrievalMode {
    /// Determine mode from query contents. Returns `None` if neither text nor
    /// vector is present (the query is empty and cannot retrieve anything).
    #[must_use]
    pub const fn determine(has_text: bool, has_vector: bool) -> Option<Self> {
        match (has_text, has_vector) {
            (true, false) => Some(Self::TextOnly),
            (false, true) => Some(Self::VectorOnly),
            (true, true) => Some(Self::Hybrid),
            (false, false) => None,
        }
    }
}

/// Route pre-retrieved result lists through the appropriate fusion path.
///
/// - **`TextOnly`**: BM25 scores cast to `f64`, order preserved.
/// - **`VectorOnly`**: ANN results converted to rank-based scores via `1.0 / (k + rank)`.
/// - **Hybrid**: delegates to [`HybridFusion::fuse`].
///
/// Both input slices must be pre-sorted "best first" by the caller (BM25
/// descending by score, ANN ascending by distance).
#[must_use]
#[allow(clippy::cast_precision_loss)] // Ranks bounded by list length, never near 2^52.
pub fn route_results(
    mode: RetrievalMode,
    bm25_results: &[(EntityId, f32)],
    ann_results: &[(EntityId, f32)],
    fusion: &HybridFusion,
) -> Vec<(EntityId, f64)> {
    match mode {
        RetrievalMode::TextOnly => bm25_results
            .iter()
            .map(|(id, score)| (*id, f64::from(*score)))
            .collect(),
        RetrievalMode::VectorOnly => {
            let k = f64::from(fusion.k);
            ann_results
                .iter()
                .enumerate()
                .map(|(i, (id, _distance))| {
                    let rank = (i + 1) as f64;
                    (*id, rrf_term(k, rank))
                })
                .collect()
        }
        RetrievalMode::Hybrid => fusion.fuse(bm25_results, ann_results),
    }
}

/// Min-max normalize a fused relevance list into a `entity_id -> score` map in
/// `[0, 1]`.
///
/// The fused scores produced by [`route_results`] are on incomparable scales
/// (raw BM25 magnitudes for `TextOnly`, tiny `1/(k+rank)` RRF terms for
/// `VectorOnly`/`Hybrid`). The ranking executor seeds its base term with a
/// normalized relevance score (`retrieval_scores`), so the SEARCH pipeline must
/// rescale the fused list to a common `[0, 1]` range first — otherwise the
/// relevance anchor's contribution would depend on which retrieval mode ran.
///
/// Normalization rules:
/// - Empty input -> empty map.
/// - A single entity, or all-equal scores (`range < EPSILON`) -> every score
///   maps to `1.0` (they are equally, maximally relevant relative to each other).
/// - Otherwise `(score - min) / (max - min)`, so the best fused result maps to
///   `1.0` and the worst to `0.0`, preserving the fused order.
///
/// Non-finite fused scores cannot occur (RRF terms and BM25 scores are finite),
/// but any that slip through are treated as the minimum and map to `0.0`.
#[must_use]
pub fn normalize_fused_scores(fused: &[(EntityId, f64)]) -> HashMap<u64, f64> {
    if fused.is_empty() {
        return HashMap::new();
    }
    let min = fused
        .iter()
        .map(|(_, s)| *s)
        .filter(|s| s.is_finite())
        .fold(f64::INFINITY, f64::min);
    let max = fused
        .iter()
        .map(|(_, s)| *s)
        .filter(|s| s.is_finite())
        .fold(f64::NEG_INFINITY, f64::max);
    let range = max - min;
    fused
        .iter()
        .map(|(id, score)| {
            let norm = if !score.is_finite() {
                0.0
            } else if range < f64::EPSILON {
                1.0
            } else {
                (score - min) / range
            };
            (id.as_u64(), norm)
        })
        .collect()
}

/// Convert ANN search results to the ranked-list format expected by fusion.
///
/// [`VectorSearchResult`] is sorted ascending by distance (best first).
/// Maps to `(EntityId, f32)` where the `f32` is the raw L2 distance,
/// preserving sort order for downstream rank computation.
#[must_use]
pub fn ann_to_ranked(ann_results: &[VectorSearchResult]) -> Vec<(EntityId, f32)> {
    ann_results
        .iter()
        .map(|r| (EntityId::new(r.id), r.distance))
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn fuse_both_lists() {
        let bm25 = vec![
            (EntityId::new(1), 1.0f32), // rank 1
            (EntityId::new(2), 0.8f32), // rank 2
            (EntityId::new(3), 0.5f32), // rank 3 (BM25 only)
        ];
        let ann = vec![
            (EntityId::new(2), 0.1f32), // rank 1 (ANN top)
            (EntityId::new(1), 0.2f32), // rank 2
            (EntityId::new(4), 0.5f32), // rank 3 (ANN only)
        ];

        let fusion = HybridFusion::new();
        let results = fusion.fuse(&bm25, &ann);

        // All four unique entities present
        let ids: Vec<u64> = results.iter().map(|(id, _)| id.as_u64()).collect();
        assert!(ids.contains(&1));
        assert!(ids.contains(&2));
        assert!(ids.contains(&3));
        assert!(ids.contains(&4));

        // Docs in both lists have higher scores than docs in one list only
        let a_score = results
            .iter()
            .find(|(id, _)| id.as_u64() == 1)
            .map(|r| r.1)
            .expect("entity 1 present");
        let c_score = results
            .iter()
            .find(|(id, _)| id.as_u64() == 3)
            .map(|r| r.1)
            .expect("entity 3 present");
        let d_score = results
            .iter()
            .find(|(id, _)| id.as_u64() == 4)
            .map(|r| r.1)
            .expect("entity 4 present");
        assert!(a_score > c_score);
        assert!(a_score > d_score);

        // Sorted descending
        let scores: Vec<f64> = results.iter().map(|(_, s)| *s).collect();
        for i in 1..scores.len() {
            assert!(scores[i - 1] >= scores[i]);
        }
    }

    #[test]
    fn fuse_bm25_only() {
        let bm25 = vec![(EntityId::new(1), 1.0f32), (EntityId::new(2), 0.5f32)];
        let fusion = HybridFusion::new();
        let results = fusion.fuse(&bm25, &[]);
        assert_eq!(results.len(), 2);
        let s1 = results
            .iter()
            .find(|(id, _)| id.as_u64() == 1)
            .map(|r| r.1)
            .expect("entity 1 present");
        let s2 = results
            .iter()
            .find(|(id, _)| id.as_u64() == 2)
            .map(|r| r.1)
            .expect("entity 2 present");
        assert!(s1 > s2);
        let expected = 1.0 / (60.0 + 1.0);
        assert!((s1 - expected).abs() < 1e-9);
    }

    #[test]
    fn fuse_ann_only() {
        let ann = vec![(EntityId::new(1), 0.1f32), (EntityId::new(2), 0.2f32)];
        let fusion = HybridFusion::new();
        let results = fusion.fuse(&[], &ann);
        assert_eq!(results.len(), 2);
        let s1 = results
            .iter()
            .find(|(id, _)| id.as_u64() == 1)
            .map(|r| r.1)
            .expect("entity 1 present");
        let s2 = results
            .iter()
            .find(|(id, _)| id.as_u64() == 2)
            .map(|r| r.1)
            .expect("entity 2 present");
        assert!(s1 > s2);
    }

    #[test]
    fn fuse_empty_lists() {
        let fusion = HybridFusion::new();
        let results = fusion.fuse(&[], &[]);
        assert!(results.is_empty());
    }

    #[test]
    fn fuse_single_doc_both_lists() {
        let bm25 = vec![(EntityId::new(1), 1.0f32)];
        let ann = vec![(EntityId::new(1), 0.1f32)];
        let fusion = HybridFusion::new();
        let results = fusion.fuse(&bm25, &ann);
        assert_eq!(results.len(), 1);
        let expected = 1.0 / (60.0 + 1.0) + 1.0 / (60.0 + 1.0);
        assert!((results[0].1 - expected).abs() < 1e-9);
    }

    /// W14 regression: equal fused scores must be ordered by ASCENDING entity id,
    /// not by randomized `HashMap` iteration order. Two items present at the same
    /// rank in both lists earn an identical RRF score, so the tie-break is the only
    /// thing that fixes their order.
    #[test]
    fn fuse_ties_break_on_ascending_entity_id() {
        // Both items rank 1 in their (single-element) lists, so each accrues the
        // SAME RRF term and the scores are exactly equal. Feed them in descending
        // id order to prove the sort actively reorders to ascending id.
        let bm25 = vec![(EntityId::new(9), 1.0f32)];
        let ann = vec![(EntityId::new(3), 0.1f32)];
        let fusion = HybridFusion::new();
        let results = fusion.fuse(&bm25, &ann);
        let ids: Vec<u64> = results.iter().map(|(id, _)| id.as_u64()).collect();
        assert_eq!(
            ids,
            vec![3, 9],
            "equal fused scores must order by ascending entity id"
        );
        // The scores really are tied (otherwise the tie-break would be untested).
        assert!((results[0].1 - results[1].1).abs() < 1e-12);
    }

    /// W14 regression: fusing identical inputs twice must yield byte-identical
    /// output, including the tie order. Before the fix the tie order depended on
    /// the per-process `HashMap` seed and could differ across runs.
    #[test]
    fn fuse_is_deterministic_across_repeated_calls() {
        let bm25: Vec<(EntityId, f32)> = (1..=20).map(|i| (EntityId::new(i), 1.0f32)).collect();
        let ann: Vec<(EntityId, f32)> = (1..=20).map(|i| (EntityId::new(i), 0.5f32)).collect();
        let fusion = HybridFusion::new();
        let first = fusion.fuse(&bm25, &ann);
        let second = fusion.fuse(&bm25, &ann);
        assert_eq!(
            first, second,
            "fusing identical inputs must produce identical, deterministic output"
        );
    }

    #[test]
    fn fuse_k_affects_scores() {
        let bm25 = vec![(EntityId::new(1), 1.0f32)];
        let ann = vec![(EntityId::new(1), 0.1f32)];
        let fusion_60 = HybridFusion::new();
        let fusion_30 = HybridFusion::with_k(30);
        let r60 = fusion_60.fuse(&bm25, &ann);
        let r30 = fusion_30.fuse(&bm25, &ann);
        // Lower k means higher individual RRF terms, so score is higher
        assert!(r30[0].1 > r60[0].1);
    }

    // --- RetrievalMode tests ---

    #[test]
    fn determine_text_only() {
        assert_eq!(
            RetrievalMode::determine(true, false),
            Some(RetrievalMode::TextOnly)
        );
    }

    #[test]
    fn determine_vector_only() {
        assert_eq!(
            RetrievalMode::determine(false, true),
            Some(RetrievalMode::VectorOnly)
        );
    }

    #[test]
    fn determine_hybrid() {
        assert_eq!(
            RetrievalMode::determine(true, true),
            Some(RetrievalMode::Hybrid)
        );
    }

    #[test]
    fn determine_none() {
        assert_eq!(RetrievalMode::determine(false, false), None);
    }

    #[test]
    fn route_text_only_passthrough() {
        let bm25 = vec![(EntityId::new(1), 1.0f32), (EntityId::new(2), 0.5f32)];
        let fusion = HybridFusion::new();
        let results = route_results(RetrievalMode::TextOnly, &bm25, &[], &fusion);
        assert_eq!(results.len(), 2);
        assert!((results[0].1 - 1.0f64).abs() < 1e-6);
        assert!((results[1].1 - 0.5f64).abs() < 1e-6);
    }

    #[test]
    fn route_vector_only_rank_based() {
        let ann = vec![(EntityId::new(1), 0.1f32), (EntityId::new(2), 0.2f32)];
        let fusion = HybridFusion::new();
        let results = route_results(RetrievalMode::VectorOnly, &[], &ann, &fusion);
        assert_eq!(results.len(), 2);
        let expected_rank1 = 1.0 / (60.0 + 1.0);
        let expected_rank2 = 1.0 / (60.0 + 2.0);
        assert!((results[0].1 - expected_rank1).abs() < 1e-9);
        assert!((results[1].1 - expected_rank2).abs() < 1e-9);
    }

    #[test]
    fn route_hybrid_calls_fuse() {
        let bm25 = vec![(EntityId::new(1), 1.0f32)];
        let ann = vec![(EntityId::new(2), 0.1f32)];
        let fusion = HybridFusion::new();
        let results = route_results(RetrievalMode::Hybrid, &bm25, &ann, &fusion);
        assert_eq!(results.len(), 2);
    }

    #[test]
    fn rrf_term_matches_inline_formula() {
        // The extracted helper must be byte-identical to the formula it replaced.
        let k = 60.0;
        for rank in 1..=10 {
            let r = f64::from(rank);
            assert!((rrf_term(k, r) - 1.0 / (k + r)).abs() < 1e-15);
        }
    }

    #[test]
    fn normalize_fused_empty_is_empty() {
        assert!(normalize_fused_scores(&[]).is_empty());
    }

    #[test]
    fn normalize_fused_single_maps_to_one() {
        let m = normalize_fused_scores(&[(EntityId::new(7), 0.0123)]);
        assert_eq!(m.len(), 1);
        assert!((m[&7] - 1.0).abs() < 1e-12);
    }

    #[test]
    fn normalize_fused_all_equal_maps_to_one() {
        let m = normalize_fused_scores(&[
            (EntityId::new(1), 0.5),
            (EntityId::new(2), 0.5),
            (EntityId::new(3), 0.5),
        ]);
        for id in 1..=3 {
            assert!((m[&id] - 1.0).abs() < 1e-12);
        }
    }

    #[test]
    fn normalize_fused_preserves_order_and_range() {
        // Descending fused scores -> best maps to 1.0, worst to 0.0, monotone.
        let m = normalize_fused_scores(&[
            (EntityId::new(10), 0.9),
            (EntityId::new(20), 0.5),
            (EntityId::new(30), 0.1),
        ]);
        assert!((m[&10] - 1.0).abs() < 1e-12);
        assert!((m[&30] - 0.0).abs() < 1e-12);
        assert!(m[&10] > m[&20] && m[&20] > m[&30]);
        for v in m.values() {
            assert!((0.0..=1.0).contains(v));
        }
    }

    #[test]
    fn ann_to_ranked_converts_correctly() {
        let ann_results = vec![
            VectorSearchResult {
                id: 42,
                distance: 0.1,
            },
            VectorSearchResult {
                id: 99,
                distance: 0.3,
            },
        ];
        let ranked = ann_to_ranked(&ann_results);
        assert_eq!(ranked.len(), 2);
        assert_eq!(ranked[0].0.as_u64(), 42);
        assert!((ranked[0].1 - 0.1f32).abs() < 1e-6);
        assert_eq!(ranked[1].0.as_u64(), 99);
    }
}

#[cfg(test)]
mod proptests {
    use proptest::prelude::*;

    use super::*;

    proptest! {
        #[test]
        fn rrf_output_is_union_of_inputs(
            bm25_ids in prop::collection::vec(1u64..=50, 0..10),
            ann_ids  in prop::collection::vec(1u64..=50, 0..10),
        ) {
            // i < 10 (proptest vec len bound) — exactly representable in f32; rank-score fixture only.
            #[allow(clippy::cast_precision_loss)]
            let bm25: Vec<(EntityId, f32)> = bm25_ids.iter().enumerate()
                .map(|(i, &id)| (EntityId::new(id), (100 - i) as f32))
                .collect();
            #[allow(clippy::cast_precision_loss)]
            let ann: Vec<(EntityId, f32)> = ann_ids.iter().enumerate()
                .map(|(i, &id)| (EntityId::new(id), i as f32 * 0.01))
                .collect();

            let fusion = HybridFusion::new();
            let results = fusion.fuse(&bm25, &ann);

            // Output is the union of unique IDs from both inputs
            let all_ids: std::collections::HashSet<u64> = bm25_ids.iter()
                .chain(ann_ids.iter())
                .copied()
                .collect();
            let result_ids: std::collections::HashSet<u64> = results.iter()
                .map(|(id, _)| id.as_u64())
                .collect();
            prop_assert_eq!(all_ids, result_ids);

            // Output is sorted descending
            for i in 1..results.len() {
                prop_assert!(results[i - 1].1 >= results[i].1);
            }
        }
    }
}