//! 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;
use crate::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 }
    }
}

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) += 1.0 / (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) += 1.0 / (k + rank);
        }

        let mut results: Vec<(EntityId, f64)> = scores
            .into_iter()
            .map(|(id, score)| (EntityId::new(id), score))
            .collect();
        results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        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, 1.0 / (k + rank))
                })
                .collect()
        }
        RetrievalMode::Hybrid => fusion.fuse(bm25_results, ann_results),
    }
}

/// 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);
    }

    #[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 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 super::*;
    use proptest::prelude::*;

    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),
        ) {
            let bm25: Vec<(EntityId, f32)> = bm25_ids.iter().enumerate()
                .map(|(i, &id)| (EntityId::new(id), (100 - i) as f32))
                .collect();
            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);
            }
        }
    }
}
