// Copyright (C) 2026  The Software Heritage developers
// See the AUTHORS file at the top-level directory of this distribution
// License: GNU General Public License version 3, or any later version
// See top-level LICENSE file for more information

use std::ops::Range;

use rapidhash::RapidHashMap;

/// Turns an iterator of `u32` into an iterator of `Range<u32>`
pub(crate) struct RangesFromItemsIter<Iter: Iterator<Item = u32>> {
    iter: Iter,
    next: Option<u32>,
}

impl<Iter: Iterator<Item = u32>> RangesFromItemsIter<Iter> {
    pub fn new(iter: impl IntoIterator<Item = u32, IntoIter = Iter>) -> Self {
        let mut iter = iter.into_iter();
        Self {
            next: iter.next(),
            iter,
        }
    }
}

impl<Iter: Iterator<Item = u32>> Iterator for RangesFromItemsIter<Iter> {
    type Item = Range<u32>;

    fn next(&mut self) -> Option<Range<u32>> {
        let next = self.next?;
        let mut range = next..next + 1;
        loop {
            let Some(next) = self.iter.next() else {
                // end of iterator
                self.next = None;
                return Some(range);
            };

            if next == range.end {
                range.end += 1;
            } else {
                // end of this run
                self.next = Some(next);
                return Some(range);
            }
        }
    }
}

#[derive(Eq, PartialEq, Hash)]
struct StringPair(Box<[u8]>, Box<[u8]>);

#[derive(Clone, Copy)]
enum LevenshteinCacheEntry {
    Exact(u32),
    LowerBound(u32),
}

#[derive(Default)]
pub struct CachedLevenshtein {
    /// maps two strings to their distance.
    ///
    /// lexicographically smallest string first in each pair
    ///
    /// levenshtein_simd_k returns `None` when a distance exceeds the higher bound we gave it, to
    /// save on computation. in that case, we cache `LevenshteinCacheEntry::LowerBound`.
    levenshtein_cache: RapidHashMap<StringPair, LevenshteinCacheEntry>,
}

impl CachedLevenshtein {
    pub fn distance<'a>(
        &mut self,
        mut first: &'a [u8],
        mut second: &'a [u8],
        max_distance: u32,
    ) -> Option<u32> {
        if first > second {
            (first, second) = (second, first);
        }
        let key = StringPair(first.into(), second.into());
        match self.levenshtein_cache.get(&key).copied() {
            Some(LevenshteinCacheEntry::Exact(distance)) => {
                if distance <= max_distance {
                    Some(distance)
                } else {
                    None
                }
            }
            Some(LevenshteinCacheEntry::LowerBound(distance_lower_bound))
                if distance_lower_bound >= max_distance =>
            {
                None
            }
            _ => {
                let distance =
                    triple_accel::levenshtein::levenshtein_simd_k(first, second, max_distance);
                self.levenshtein_cache.insert(
                    key,
                    match distance {
                        Some(distance) => LevenshteinCacheEntry::Exact(distance),
                        None => LevenshteinCacheEntry::LowerBound(max_distance),
                    },
                );
                if self.levenshtein_cache.len() >= 1_000_000 {
                    // evict 25% of the keys at "random".
                    // We don't use a LRU because callers compute distances always in the same
                    // order, so if there are more lines than the size of the cache, an LRU
                    // algorithm would cause 100% cache miss.
                    let mut i = 0;
                    self.levenshtein_cache.retain(|_k, _v| {
                        i += 1;
                        i % 4 != 0
                    });
                }
                distance
            }
        }
    }
}

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

    #[test]
    fn test_rangesfromitemsiter() {
        assert_eq!(
            RangesFromItemsIter::new(vec![3, 5, 6, 7, 10, 11, 13, 15]).collect::<Vec<_>>(),
            vec![3..4, 5..8, 10..12, 13..14, 15..16],
        );
    }

    #[test]
    fn test_cached_levenshtein_distance() {
        let mut cache = CachedLevenshtein::default();
        assert_eq!(cache.distance(b"foo", b"bar", 100), Some(3));
        assert_eq!(cache.distance(b"qux", b"quux", 100), Some(1));
        assert_eq!(cache.distance(b"", b"", 100), Some(0));
        assert_eq!(cache.distance(b"foo", b"foo", 100), Some(0));
        assert_eq!(cache.distance(b"", b"foo", 100), Some(3));
        assert_eq!(cache.distance(b"foo", b"", 100), Some(3));
    }

    #[test]
    fn test_cached_levenshtein_max_distance() {
        let mut cache = CachedLevenshtein::default();
        assert_eq!(cache.distance(b"foo", b"bar", 3), Some(3));
        assert_eq!(cache.distance(b"foo", b"bar", 2), None);
        assert_eq!(cache.distance(b"foo", b"bar", 0), None);
    }

    #[test]
    fn test_cached_levenshtein_repeated_query() {
        let mut cache = CachedLevenshtein::default();
        assert_eq!(cache.distance(b"foo", b"bar", 100), Some(3));
        assert_eq!(cache.distance(b"foo", b"bar", 100), Some(3));
        assert_eq!(cache.distance(b"foo", b"bar", 2), None);
        assert_eq!(cache.distance(b"foo", b"bar", 3), Some(3));
    }

    #[test]
    fn test_cached_levenshtein_lowerbound_then_recompute() {
        let mut cache = CachedLevenshtein::default();
        assert_eq!(cache.distance(b"foo", b"bar", 2), None);
        assert_eq!(cache.distance(b"foo", b"bar", 1), None);
        assert_eq!(cache.distance(b"foo", b"bar", 2), None);
        assert_eq!(cache.distance(b"foo", b"bar", 3), Some(3));
        assert_eq!(cache.distance(b"foo", b"bar", 3), Some(3));
        assert_eq!(cache.distance(b"foo", b"bar", 100), Some(3));
    }

    #[test]
    fn test_cached_levenshtein_lowerbound_then_recompute2() {
        let mut cache = CachedLevenshtein::default();
        assert_eq!(cache.distance(b"foo", b"bar", 1), None);
        assert_eq!(cache.distance(b"foo", b"bar", 2), None);
        assert_eq!(cache.distance(b"foo", b"bar", 3), Some(3));
        assert_eq!(cache.distance(b"foo", b"bar", 3), Some(3));
        assert_eq!(cache.distance(b"foo", b"bar", 100), Some(3));
    }
}
