//! B-tree backed range index for ordered numeric fields.
//!
//! Maps attribute values to `RoaringBitmap` sets of entity IDs.
//! The B-tree ordering enables efficient range queries.

use std::collections::BTreeMap;
use std::ops::Bound;
use std::sync::RwLock;

use roaring::RoaringBitmap;

use crate::schema::EntityId;
use crate::storage::keys::{Tag, encode_key};

use super::IndexError;

use crate::storage::indexes::bitmap::INDEX_ROOT_ID;

/// Serialized key-value pairs for index persistence.
type KvPairs = Vec<(Vec<u8>, Vec<u8>)>;

/// A B-tree backed range index for a single ordered numeric field.
///
/// Maps attribute values to `RoaringBitmap` sets of entity IDs.
/// The B-tree ordering enables range queries: "all entities with duration >= 300"
/// unions bitmaps for keys 300, 301, ...
///
/// # Concurrency
///
/// Same model as `BitmapIndex`: `RwLock` for read/write separation.
pub struct RangeIndex<V: Ord + Clone> {
    field_name: String,
    tree: RwLock<BTreeMap<V, RoaringBitmap>>,
}

impl<V: Ord + Clone> RangeIndex<V> {
    /// Create a new, empty range index for the given field.
    #[must_use]
    pub fn new(field_name: impl Into<String>) -> Self {
        Self {
            field_name: field_name.into(),
            tree: RwLock::new(BTreeMap::new()),
        }
    }

    /// The field name this index covers.
    #[must_use]
    pub fn field_name(&self) -> &str {
        &self.field_name
    }

    /// Add an entity with the given attribute value.
    pub fn insert(&self, entity_id: u32, value: V) {
        let mut tree = self
            .tree
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let bitmap = tree.entry(value).or_default();
        bitmap.insert(entity_id);
        drop(tree);
    }

    /// Remove an entity from the bitmap at the given value.
    ///
    /// Returns `true` if the entity was present and removed.
    pub fn delete(&self, entity_id: u32, value: &V) -> bool {
        let mut tree = self
            .tree
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let removed = if let Some(bitmap) = tree.get_mut(value) {
            bitmap.remove(entity_id)
        } else {
            return false;
        };
        // Clean up empty bitmaps.
        if tree.get(value).is_some_and(RoaringBitmap::is_empty) {
            tree.remove(value);
        }
        removed
    }

    /// Range query: union of all bitmaps with keys in [lo, hi].
    #[must_use]
    pub fn range(&self, lo: Bound<&V>, hi: Bound<&V>) -> RoaringBitmap {
        let tree = self
            .tree
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let mut result = RoaringBitmap::new();
        for (_key, bitmap) in tree.range((lo, hi)) {
            result |= bitmap;
        }
        drop(tree);
        result
    }

    /// Greater-than query: entities with value > threshold.
    #[must_use]
    pub fn gt(&self, threshold: &V) -> RoaringBitmap {
        self.range(Bound::Excluded(threshold), Bound::Unbounded)
    }

    /// Greater-than-or-equal query.
    #[must_use]
    pub fn gte(&self, threshold: &V) -> RoaringBitmap {
        self.range(Bound::Included(threshold), Bound::Unbounded)
    }

    /// Less-than query: entities with value < threshold.
    #[must_use]
    pub fn lt(&self, threshold: &V) -> RoaringBitmap {
        self.range(Bound::Unbounded, Bound::Excluded(threshold))
    }

    /// Less-than-or-equal query.
    #[must_use]
    pub fn lte(&self, threshold: &V) -> RoaringBitmap {
        self.range(Bound::Unbounded, Bound::Included(threshold))
    }

    /// Estimate fraction of entities in range. Returns value in [0.0, 1.0].
    ///
    /// Returns 0.0 if `total` is 0.
    #[must_use]
    pub fn selectivity(&self, lo: Bound<&V>, hi: Bound<&V>, total: u64) -> f64 {
        if total == 0 {
            return 0.0;
        }
        let range_bitmap = self.range(lo, hi);
        #[allow(clippy::cast_precision_loss)]
        let result = range_bitmap.len() as f64 / total as f64;
        result.clamp(0.0, 1.0)
    }

    /// Total distinct entity IDs indexed.
    #[must_use]
    pub fn total_count(&self) -> u64 {
        let tree = self
            .tree
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let mut union = RoaringBitmap::new();
        for bitmap in tree.values() {
            union |= bitmap;
        }
        drop(tree);
        union.len()
    }

    /// Number of distinct attribute values in the tree.
    #[must_use]
    pub fn distinct_values(&self) -> usize {
        let tree = self
            .tree
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        tree.len()
    }

    /// Whether the index is empty.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        let tree = self
            .tree
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        tree.is_empty()
    }
}

// `RwLock<BTreeMap<V, RoaringBitmap>>` is `Send + Sync` when `V: Send + Sync`.
// The compiler derives these automatically; no manual unsafe impl is needed.

// === Persistence for concrete types ===

impl RangeIndex<u64> {
    /// Serialize all bitmaps to (key, value) pairs for storage.
    ///
    /// Key suffix: `b"RNG:{field_name}:{value_be_bytes}"`.
    ///
    /// # Errors
    ///
    /// Returns `IndexError::Serialization` if serialization fails.
    pub fn serialize_to_kv_pairs(&self) -> Result<KvPairs, IndexError> {
        let tree = self
            .tree
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let mut pairs = Vec::with_capacity(tree.len());
        for (value, bitmap) in tree.iter() {
            let mut suffix = format!("RNG:{}:", self.field_name).into_bytes();
            suffix.extend_from_slice(&value.to_be_bytes());
            let key = encode_key(EntityId::new(INDEX_ROOT_ID), Tag::Idx, &suffix);
            let mut bytes = Vec::new();
            bitmap
                .serialize_into(&mut bytes)
                .map_err(|e| IndexError::Serialization(e.to_string()))?;
            pairs.push((key, bytes));
        }
        drop(tree);
        Ok(pairs)
    }

    /// Deserialize from (key, value) pairs.
    ///
    /// # Errors
    ///
    /// Returns `IndexError::Serialization` if deserialization fails.
    pub fn load_from_kv_pairs(
        field_name: impl Into<String>,
        pairs: impl Iterator<Item = (Vec<u8>, Vec<u8>)>,
    ) -> Result<Self, IndexError> {
        let field_name = field_name.into();
        let prefix_bytes = format!("RNG:{field_name}:").into_bytes();
        let mut tree = BTreeMap::new();

        for (key, value_bytes) in pairs {
            if key.len() < 10 {
                continue;
            }
            let suffix_raw = &key[10..];
            if !suffix_raw.starts_with(&prefix_bytes) {
                continue;
            }
            let value_raw = &suffix_raw[prefix_bytes.len()..];
            if value_raw.len() < 8 {
                return Err(IndexError::Serialization(
                    "truncated u64 value in key".to_string(),
                ));
            }
            let value = u64::from_be_bytes(value_raw[..8].try_into().map_err(
                |e: std::array::TryFromSliceError| IndexError::Serialization(e.to_string()),
            )?);
            let bitmap = RoaringBitmap::deserialize_from(value_bytes.as_slice())
                .map_err(|e| IndexError::Serialization(e.to_string()))?;
            if !bitmap.is_empty() {
                tree.insert(value, bitmap);
            }
        }

        Ok(Self {
            field_name,
            tree: RwLock::new(tree),
        })
    }
}

impl RangeIndex<u32> {
    /// Serialize all bitmaps to (key, value) pairs for storage.
    ///
    /// # Errors
    ///
    /// Returns `IndexError::Serialization` if serialization fails.
    pub fn serialize_to_kv_pairs(&self) -> Result<KvPairs, IndexError> {
        let tree = self
            .tree
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let mut pairs = Vec::with_capacity(tree.len());
        for (value, bitmap) in tree.iter() {
            let mut suffix = format!("RNG:{}:", self.field_name).into_bytes();
            suffix.extend_from_slice(&value.to_be_bytes());
            let key = encode_key(EntityId::new(INDEX_ROOT_ID), Tag::Idx, &suffix);
            let mut bytes = Vec::new();
            bitmap
                .serialize_into(&mut bytes)
                .map_err(|e| IndexError::Serialization(e.to_string()))?;
            pairs.push((key, bytes));
        }
        drop(tree);
        Ok(pairs)
    }

    /// Deserialize from (key, value) pairs.
    ///
    /// # Errors
    ///
    /// Returns `IndexError::Serialization` if deserialization fails.
    pub fn load_from_kv_pairs(
        field_name: impl Into<String>,
        pairs: impl Iterator<Item = (Vec<u8>, Vec<u8>)>,
    ) -> Result<Self, IndexError> {
        let field_name = field_name.into();
        let prefix = format!("RNG:{field_name}:");
        let mut tree = BTreeMap::new();

        for (key, value_bytes) in pairs {
            if key.len() < 10 {
                continue;
            }
            let suffix_raw = &key[10..];
            let prefix_bytes = prefix.as_bytes();
            if !suffix_raw.starts_with(prefix_bytes) {
                continue;
            }
            let value_bytes_raw = &suffix_raw[prefix_bytes.len()..];
            if value_bytes_raw.len() < 4 {
                return Err(IndexError::Serialization(
                    "truncated u32 value in key".to_string(),
                ));
            }
            let value = u32::from_be_bytes(value_bytes_raw[..4].try_into().map_err(
                |e: std::array::TryFromSliceError| IndexError::Serialization(e.to_string()),
            )?);
            let bitmap = RoaringBitmap::deserialize_from(value_bytes.as_slice())
                .map_err(|e| IndexError::Serialization(e.to_string()))?;
            if !bitmap.is_empty() {
                tree.insert(value, bitmap);
            }
        }

        Ok(Self {
            field_name,
            tree: RwLock::new(tree),
        })
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::float_cmp)]
mod tests {
    use super::*;
    use proptest::prelude::*;

    // === Unit tests ===

    #[test]
    fn new_index_is_empty() {
        let index: RangeIndex<u32> = RangeIndex::new("duration");
        assert!(index.is_empty());
        assert_eq!(index.total_count(), 0);
        assert_eq!(index.distinct_values(), 0);
    }

    #[test]
    fn insert_and_range_query() {
        let index: RangeIndex<u32> = RangeIndex::new("duration");
        index.insert(1, 60);
        index.insert(2, 300);
        index.insert(3, 600);
        index.insert(4, 1800);

        let result = index.range(Bound::Included(&300), Bound::Included(&600));
        assert_eq!(result.len(), 2);
        assert!(result.contains(2));
        assert!(result.contains(3));
        assert!(!result.contains(1));
        assert!(!result.contains(4));
    }

    #[test]
    fn gte_query() {
        let index: RangeIndex<u32> = RangeIndex::new("duration");
        index.insert(1, 60);
        index.insert(2, 300);
        index.insert(3, 600);

        let result = index.gte(&300);
        assert_eq!(result.len(), 2);
        assert!(result.contains(2));
        assert!(result.contains(3));
    }

    #[test]
    fn gt_query() {
        let index: RangeIndex<u32> = RangeIndex::new("duration");
        index.insert(1, 60);
        index.insert(2, 300);
        index.insert(3, 600);

        let result = index.gt(&300);
        assert_eq!(result.len(), 1);
        assert!(result.contains(3));
    }

    #[test]
    fn lt_and_lte_queries() {
        let index: RangeIndex<u32> = RangeIndex::new("duration");
        index.insert(1, 60);
        index.insert(2, 300);
        index.insert(3, 600);

        let lt = index.lt(&300);
        assert_eq!(lt.len(), 1);
        assert!(lt.contains(1));

        let lte = index.lte(&300);
        assert_eq!(lte.len(), 2);
        assert!(lte.contains(1));
        assert!(lte.contains(2));
    }

    #[test]
    fn unbounded_range_returns_all() {
        let index: RangeIndex<u32> = RangeIndex::new("duration");
        index.insert(1, 60);
        index.insert(2, 300);
        index.insert(3, 600);

        let all = index.range(Bound::Unbounded, Bound::Unbounded);
        assert_eq!(all.len(), 3);
    }

    #[test]
    fn empty_range_returns_empty_bitmap() {
        let index: RangeIndex<u32> = RangeIndex::new("duration");
        index.insert(1, 60);
        index.insert(2, 300);

        let result = index.range(Bound::Included(&400), Bound::Included(&500));
        assert!(result.is_empty());
    }

    #[test]
    fn selectivity_full_range() {
        let index: RangeIndex<u32> = RangeIndex::new("duration");
        index.insert(1, 60);
        index.insert(2, 300);
        index.insert(3, 600);

        let sel = index.selectivity(Bound::Unbounded, Bound::Unbounded, 3);
        assert!((sel - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn selectivity_partial_range() {
        let index: RangeIndex<u32> = RangeIndex::new("duration");
        index.insert(1, 60);
        index.insert(2, 300);
        index.insert(3, 600);
        index.insert(4, 1800);

        let sel = index.selectivity(Bound::Included(&300), Bound::Included(&600), 4);
        assert!((sel - 0.5).abs() < f64::EPSILON);
    }

    #[test]
    fn selectivity_zero_total() {
        let index: RangeIndex<u32> = RangeIndex::new("duration");
        let sel = index.selectivity(Bound::Unbounded, Bound::Unbounded, 0);
        assert!((sel - 0.0).abs() < f64::EPSILON);
    }

    #[test]
    fn delete_cleans_up_empty_entries() {
        let index: RangeIndex<u32> = RangeIndex::new("duration");
        index.insert(1, 300);
        index.delete(1, &300);

        assert_eq!(index.total_count(), 0);
        assert_eq!(index.distinct_values(), 0);
    }

    #[test]
    fn multiple_entities_same_value() {
        let index: RangeIndex<u32> = RangeIndex::new("duration");
        index.insert(1, 300);
        index.insert(2, 300);
        index.insert(3, 300);

        let result = index.gte(&300);
        assert_eq!(result.len(), 3);
        assert_eq!(index.distinct_values(), 1);
    }

    #[test]
    fn timestamp_index_u64() {
        let index: RangeIndex<u64> = RangeIndex::new("created_at");
        let now_ns: u64 = 1_708_000_000_000_000_000;
        let one_day_ago = now_ns - 86_400_000_000_000;
        let seven_days_ago = now_ns - 7 * 86_400_000_000_000;

        index.insert(1, now_ns);
        index.insert(2, one_day_ago);
        index.insert(3, seven_days_ago);
        index.insert(4, seven_days_ago - 1);

        let recent = index.gte(&seven_days_ago);
        assert_eq!(recent.len(), 3);
        assert!(!recent.contains(4));
    }

    // === Property tests ===

    proptest! {
        #[test]
        fn range_query_correctness(
            entries in prop::collection::vec(
                (0u32..10_000, 0u32..1000),
                1..200,
            ),
            lo in 0u32..500,
            hi in 500u32..1000,
        ) {
            let index: RangeIndex<u32> = RangeIndex::new("test_field");
            for &(id, value) in &entries {
                index.insert(id, value);
            }

            let result = index.range(Bound::Included(&lo), Bound::Included(&hi));

            for &(id, value) in &entries {
                if value >= lo && value <= hi {
                    prop_assert!(result.contains(id),
                        "entity {id} with value {value} should be in range [{lo}, {hi}]");
                }
            }
        }

        #[test]
        fn selectivity_in_unit_range(
            entries in prop::collection::vec(
                (0u32..10_000, 0u32..1000),
                1..200,
            ),
            lo in 0u32..500,
            hi in 500u32..1000,
        ) {
            let index: RangeIndex<u32> = RangeIndex::new("test_field");
            for &(id, value) in &entries {
                index.insert(id, value);
            }
            let total = index.total_count();
            let sel = index.selectivity(Bound::Included(&lo), Bound::Included(&hi), total);
            prop_assert!(sel >= 0.0, "selectivity was {sel}");
            prop_assert!(sel <= 1.0, "selectivity was {sel}");
        }

        #[test]
        fn insert_delete_roundtrip(
            entries in prop::collection::vec(
                (0u32..1_000, 0u32..100),
                1..100,
            ),
        ) {
            let index: RangeIndex<u32> = RangeIndex::new("test_field");
            for &(id, value) in &entries {
                index.insert(id, value);
            }

            for &(id, value) in &entries {
                index.delete(id, &value);
            }

            let result = index.range(Bound::Unbounded, Bound::Unbounded);
            prop_assert!(result.is_empty(), "expected empty after deleting all");
        }

        #[test]
        fn serialize_roundtrip_u32(
            entries in prop::collection::vec(
                (0u32..10_000, 0u32..1000),
                1..100,
            ),
        ) {
            let index: RangeIndex<u32> = RangeIndex::new("duration");
            for &(id, value) in &entries {
                index.insert(id, value);
            }

            let kv_pairs = index.serialize_to_kv_pairs().unwrap();
            let restored = RangeIndex::<u32>::load_from_kv_pairs(
                "duration",
                kv_pairs.into_iter(),
            ).unwrap();

            let orig = index.range(Bound::Unbounded, Bound::Unbounded);
            let rest = restored.range(Bound::Unbounded, Bound::Unbounded);
            prop_assert_eq!(orig, rest);
        }

        #[test]
        fn serialize_roundtrip_u64(
            entries in prop::collection::vec((any::<u32>(), any::<u64>()), 1..50),
        ) {
            let index: RangeIndex<u64> = RangeIndex::new("created_at");
            for &(entity_id, ts) in &entries {
                index.insert(entity_id, ts);
            }
            let pairs = index.serialize_to_kv_pairs().unwrap();
            let restored = RangeIndex::<u64>::load_from_kv_pairs("created_at", pairs.into_iter()).unwrap();

            // All values in original must be present in restored.
            for &(entity_id, ts) in &entries {
                let orig_bm = index.range(
                    Bound::Included(&ts),
                    Bound::Included(&ts),
                );
                let rest_bm = restored.range(
                    Bound::Included(&ts),
                    Bound::Included(&ts),
                );
                prop_assert!(orig_bm.contains(entity_id), "entity missing in original");
                prop_assert!(rest_bm.contains(entity_id), "entity missing in restored after roundtrip");
            }
            prop_assert_eq!(index.total_count(), restored.total_count());
        }
    }
}
