//! Hot/cold merge algorithm for vector search results.
//!
//! When searching across hot (in-memory) and cold (disk-backed) indexes,
//! results must be merged while:
//! - Maintaining distance ordering (ascending)
//! - Deduplicating by hash (preferring hot over cold)
//! - Respecting the k limit

use std::collections::HashSet;
use stemedb_core::types::Hash;

/// Merge search results from hot and cold indexes.
///
/// Both input vectors must be sorted by distance (ascending).
/// Deduplicates by hash, preferring hot results over cold.
///
/// # Arguments
/// * `hot_results` - Results from the hot (recent) index, sorted by distance
/// * `cold_results` - Results from the cold (historical) index, sorted by distance
/// * `k` - Maximum number of results to return
///
/// # Returns
/// Merged results sorted by distance, limited to k items.
///
/// # Algorithm
///
/// Uses a two-pointer interleave merge:
/// 1. Track which hashes we've seen (to deduplicate)
/// 2. Compare distances from hot and cold heads
/// 3. Take the smaller distance, skip if hash already seen
/// 4. Stop when we have k results or exhaust both inputs
///
/// Complexity: O(hot.len() + cold.len()) time, O(k) space for result
pub fn merge_search_results(
    hot_results: Vec<(Hash, f32)>,
    cold_results: Vec<(Hash, f32)>,
    k: usize,
) -> Vec<(Hash, f32)> {
    if k == 0 {
        return Vec::new();
    }

    let mut result = Vec::with_capacity(k);
    let mut seen: HashSet<Hash> = HashSet::with_capacity(k);

    let mut hot_idx = 0;
    let mut cold_idx = 0;

    while result.len() < k && (hot_idx < hot_results.len() || cold_idx < cold_results.len()) {
        // Determine which source has the smaller distance
        let take_hot = match (hot_results.get(hot_idx), cold_results.get(cold_idx)) {
            (Some((_, hot_dist)), Some((_, cold_dist))) => hot_dist <= cold_dist,
            (Some(_), None) => true,
            (None, Some(_)) => false,
            (None, None) => break,
        };

        if take_hot {
            let (hash, distance) = hot_results[hot_idx];
            hot_idx += 1;

            // Always include hot results (they take precedence for deduplication)
            if seen.insert(hash) {
                result.push((hash, distance));
            }
        } else {
            let (hash, distance) = cold_results[cold_idx];
            cold_idx += 1;

            // Only include cold results if not already seen from hot
            if seen.insert(hash) {
                result.push((hash, distance));
            }
        }
    }

    result
}

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

    fn hash(n: u8) -> Hash {
        let mut h = [0u8; 32];
        h[0] = n;
        h
    }

    #[test]
    fn test_merge_empty_inputs() {
        let result = merge_search_results(vec![], vec![], 5);
        assert!(result.is_empty());
    }

    #[test]
    fn test_merge_hot_only() {
        let hot = vec![(hash(1), 0.1), (hash(2), 0.2), (hash(3), 0.3)];
        let cold = vec![];

        let result = merge_search_results(hot, cold, 5);
        assert_eq!(result.len(), 3);
        assert_eq!(result[0].0, hash(1));
        assert_eq!(result[1].0, hash(2));
        assert_eq!(result[2].0, hash(3));
    }

    #[test]
    fn test_merge_cold_only() {
        let hot = vec![];
        let cold = vec![(hash(1), 0.1), (hash(2), 0.2)];

        let result = merge_search_results(hot, cold, 5);
        assert_eq!(result.len(), 2);
        assert_eq!(result[0].0, hash(1));
        assert_eq!(result[1].0, hash(2));
    }

    #[test]
    fn test_merge_interleaved() {
        // Hot: distances 0.1, 0.3, 0.5
        // Cold: distances 0.2, 0.4, 0.6
        let hot = vec![(hash(1), 0.1), (hash(3), 0.3), (hash(5), 0.5)];
        let cold = vec![(hash(2), 0.2), (hash(4), 0.4), (hash(6), 0.6)];

        let result = merge_search_results(hot, cold, 6);
        assert_eq!(result.len(), 6);

        // Should be sorted by distance
        assert_eq!(result[0].0, hash(1)); // 0.1
        assert_eq!(result[1].0, hash(2)); // 0.2
        assert_eq!(result[2].0, hash(3)); // 0.3
        assert_eq!(result[3].0, hash(4)); // 0.4
        assert_eq!(result[4].0, hash(5)); // 0.5
        assert_eq!(result[5].0, hash(6)); // 0.6
    }

    #[test]
    fn test_merge_respects_k_limit() {
        let hot = vec![(hash(1), 0.1), (hash(2), 0.2), (hash(3), 0.3)];
        let cold = vec![(hash(4), 0.4), (hash(5), 0.5), (hash(6), 0.6)];

        let result = merge_search_results(hot, cold, 3);
        assert_eq!(result.len(), 3);
        assert_eq!(result[0].0, hash(1));
        assert_eq!(result[1].0, hash(2));
        assert_eq!(result[2].0, hash(3));
    }

    #[test]
    fn test_merge_deduplicates_by_distance_order() {
        // Same hash appears in both - the first occurrence by distance wins
        // Cold has hash(1) at distance 0.1, hot has hash(1) at distance 0.5
        // Since we merge by distance order, cold(1) at 0.1 comes first and wins
        let hot = vec![(hash(1), 0.5)]; // hash(1) at distance 0.5
        let cold = vec![(hash(1), 0.1), (hash(2), 0.2)]; // hash(1) at distance 0.1 (better!)

        let result = merge_search_results(hot, cold, 5);

        // hash(1) from cold comes first (0.1)
        // hash(2) from cold comes second (0.2)
        // hash(1) from hot (0.5) is skipped because hash(1) already seen
        assert_eq!(result.len(), 2);
        assert_eq!(result[0].0, hash(1)); // 0.1 from cold
        assert_eq!(result[1].0, hash(2)); // 0.2 from cold
    }

    #[test]
    fn test_merge_k_zero() {
        let hot = vec![(hash(1), 0.1)];
        let cold = vec![(hash(2), 0.2)];

        let result = merge_search_results(hot, cold, 0);
        assert!(result.is_empty());
    }

    #[test]
    fn test_merge_deduplicates_cold_duplicates() {
        // Multiple entries with same hash should be deduplicated
        let hot = vec![];
        let cold = vec![(hash(1), 0.1), (hash(1), 0.2), (hash(2), 0.3)];

        let result = merge_search_results(hot, cold, 5);
        assert_eq!(result.len(), 2);
        assert_eq!(result[0].0, hash(1)); // 0.1 (first occurrence)
        assert_eq!(result[1].0, hash(2)); // 0.3
    }

    #[test]
    fn test_merge_equal_distances_prefers_hot() {
        // When distances are equal, take from hot first
        let hot = vec![(hash(1), 0.5)];
        let cold = vec![(hash(2), 0.5)];

        let result = merge_search_results(hot, cold, 2);
        assert_eq!(result.len(), 2);
        // Hot should come first when distances are equal
        assert_eq!(result[0].0, hash(1));
        assert_eq!(result[1].0, hash(2));
    }
}
