//! Roaring bitmap index for categorical metadata fields.
//!
//! Maps (`field_name`, `field_value`) pairs to `RoaringBitmap` sets of entity IDs.
//! Used for exact-match filter predicates: `category:jazz`, `format:video`, etc.

use std::collections::HashMap;
use std::sync::{Arc, RwLock};

use roaring::RoaringBitmap;

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

use super::IndexError;

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

/// Reserved entity ID anchoring all index storage keys.
///
/// All index metadata is stored under `EntityId(0)` with `Tag::Idx`, which
/// guarantees no collision with legitimate entity data:
/// - Entity 0 metadata uses `Tag::Meta`
/// - Entity 0 signals use `Tag::Signal`
/// - Index data uses `Tag::Idx` exclusively
///
/// This allows efficient prefix scans over all index keys via
/// `entity_tag_prefix(EntityId(0), Tag::Idx)`.
pub const INDEX_ROOT_ID: u64 = 0;

/// A roaring bitmap index for a single categorical metadata field.
///
/// Maps field values (strings) to `RoaringBitmap` sets of entity IDs.
/// Supports exact-match lookups, multi-value fields (tags), and cardinality
/// queries for selectivity estimation.
///
/// # Concurrency
///
/// Reads and writes are protected by an `RwLock`. At M2 scale (10K entities),
/// contention is negligible. Reads take the read lock; writes take the write lock.
///
/// # Persistence
///
/// Each (`field_name`, `field_value`) bitmap is stored as:
/// ```text
/// Key: [INDEX_ROOT_ID: 8 bytes BE][0x00][Tag::Idx][b"BMP:"][field_name][b":"][field_value]
/// Value: RoaringBitmap serialized bytes
/// ```
pub struct BitmapIndex {
    field_name: String,
    values: Arc<RwLock<HashMap<String, RoaringBitmap>>>,
}

impl Clone for BitmapIndex {
    fn clone(&self) -> Self {
        Self {
            field_name: self.field_name.clone(),
            values: Arc::clone(&self.values),
        }
    }
}

impl BitmapIndex {
    /// Create a new, empty bitmap index for the given field.
    #[must_use]
    pub fn new(field_name: impl Into<String>) -> Self {
        Self {
            field_name: field_name.into(),
            values: Arc::new(RwLock::new(HashMap::new())),
        }
    }

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

    /// Add an entity to the bitmap for the given field value.
    ///
    /// Returns `true` if the entity was newly added (was not already present).
    pub fn insert(&self, entity_id: u32, value: impl Into<String>) -> bool {
        let mut map = self
            .values
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let bitmap = map.entry(value.into()).or_default();
        let inserted = bitmap.insert(entity_id);
        drop(map);
        inserted
    }

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

    /// Remove an entity from ALL value bitmaps.
    ///
    /// Returns the number of bitmaps the entity was removed from.
    pub fn delete_entity(&self, entity_id: u32) -> usize {
        let mut map = self
            .values
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let mut count = 0;
        let mut empty_keys = Vec::new();
        for (key, bitmap) in map.iter_mut() {
            if bitmap.remove(entity_id) {
                count += 1;
                if bitmap.is_empty() {
                    empty_keys.push(key.clone());
                }
            }
        }
        for key in empty_keys {
            map.remove(&key);
        }
        count
    }

    /// Look up the bitmap for an exact field value. Returns a clone.
    ///
    /// Returns `None` if no entities have this value.
    #[must_use]
    pub fn get(&self, value: &str) -> Option<RoaringBitmap> {
        let map = self
            .values
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        map.get(value).filter(|b| !b.is_empty()).cloned()
    }

    /// Look up bitmaps for multiple values and return their union.
    ///
    /// Implements OR-within-dimension: `category IN [jazz, blues]`.
    #[must_use]
    pub fn get_union(&self, values: &[&str]) -> RoaringBitmap {
        let map = self
            .values
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let mut result = RoaringBitmap::new();
        for v in values {
            if let Some(bitmap) = map.get(*v) {
                result |= bitmap;
            }
        }
        result
    }

    /// Number of entities matching a specific field value.
    #[must_use]
    pub fn cardinality(&self, value: &str) -> u64 {
        let map = self
            .values
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        map.get(value).map_or(0, RoaringBitmap::len)
    }

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

    /// Enumerate all indexed field values.
    #[must_use]
    pub fn values(&self) -> Vec<String> {
        let map = self
            .values
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        map.keys().cloned().collect()
    }

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

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

    /// Total number of entity IDs across all bitmap entries (sum, not union).
    ///
    /// Unlike [`total_count`](Self::total_count) which de-duplicates via union,
    /// this returns the raw sum of per-value bitmap cardinalities. Useful for
    /// monitoring index size without the cost of computing a full union.
    #[must_use]
    pub fn total_cardinality(&self) -> u64 {
        let map = self
            .values
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        map.values().map(RoaringBitmap::len).sum()
    }

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

    /// Deserialize bitmaps from (key, value) pairs.
    ///
    /// Scans for keys with the `BMP:{field_name}:` prefix and deserializes.
    ///
    /// # 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!("BMP:{field_name}:");
        let mut map = HashMap::new();

        for (key, value_bytes) in pairs {
            // Parse suffix from encoded key (skip 10-byte prefix: 8+1+1)
            if key.len() < 10 {
                continue;
            }
            let suffix = &key[10..];
            let suffix_str = std::str::from_utf8(suffix)
                .map_err(|e| IndexError::Serialization(e.to_string()))?;
            if let Some(field_value) = suffix_str.strip_prefix(&prefix) {
                let bitmap = RoaringBitmap::deserialize_from(value_bytes.as_slice())
                    .map_err(|e| IndexError::Serialization(e.to_string()))?;
                if !bitmap.is_empty() {
                    map.insert(field_value.to_string(), bitmap);
                }
            }
        }

        Ok(Self {
            field_name,
            values: Arc::new(RwLock::new(map)),
        })
    }
}

// `Arc<RwLock<HashMap<String, RoaringBitmap>>>` is `Send + Sync` automatically
// because `RoaringBitmap`, `String`, `RwLock`, and `Arc` are all `Send + Sync`.
// No manual unsafe impl needed.
#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests;
