//! 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,
    sync::{Arc, RwLock},
};

use roaring::RoaringBitmap;

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

/// 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`. The tag
/// byte is what isolates index data from entity data at the same ID, so
/// `EntityId(0)`'s own rows under other tags (`Tag::Meta`, `Tag::Sig`, etc.)
/// never collide with index rows:
/// - Entity 0 metadata uses `Tag::Meta`
/// - Entity 0 signal state uses `Tag::Sig`
/// - 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_len: u16 BE][field_name][field_value]
/// Value: RoaringBitmap serialized bytes
/// ```
///
/// The field name is **length-prefixed** (a big-endian `u16` before the field
/// bytes) so the field/value boundary is unambiguous. A naive `BMP:{field}:{value}`
/// delimiter scheme collides across indexes whenever a field name (or value)
/// contains the delimiter -- e.g. field `"a:b"` + value `"c"` and field `"a"` +
/// value `"b:c"` both render as `BMP:a:b:c`. Length-prefixing removes that
/// ambiguity entirely, which matters once arbitrary user-defined field names
/// are indexable.
#[derive(Debug)]
pub struct BitmapIndex {
    field_name: String,
    values: Arc<RwLock<HashMap<String, RoaringBitmap>>>,
    /// Cached distinct-entity union cardinality, invalidated on every mutation.
    ///
    /// `total_count()` unions every value bitmap — `O(distinct_values *
    /// avg_bitmap_size)`. Memoizing it makes the planner's repeated
    /// per-predicate reads between writes O(1). Shared via `Arc` so a cloned
    /// handle (which shares `values`) also shares the memo. It is a read-side
    /// memo of state already in `values`, never a source of truth: a missed
    /// invalidation would only force a recompute, never return a wrong count.
    total_count_cache: Arc<RwLock<Option<u64>>>,
}

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

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())),
            total_count_cache: Arc::new(RwLock::new(Some(0))),
        }
    }

    /// Invalidate the memoized total-count union after a mutation.
    fn invalidate_total_count(&self) {
        *self
            .total_count_cache
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner) = None;
    }

    /// 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);
        if inserted {
            self.invalidate_total_count();
        }
        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);
        // Capture both the removal result and whether the bitmap is now empty in
        // the single `get_mut` borrow, so the cleanup branch below needs no
        // second hash lookup.
        let Some(bitmap) = map.get_mut(value) else {
            return false;
        };
        let removed = bitmap.remove(entity_id);
        let now_empty = bitmap.is_empty();
        // Clean up empty bitmaps to avoid dead entries.
        if now_empty {
            map.remove(value);
        }
        drop(map);
        if removed {
            self.invalidate_total_count();
        }
        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);
        }
        drop(map);
        if count > 0 {
            self.invalidate_total_count();
        }
        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).
    ///
    /// Memoized: the union over every value bitmap is computed at most once per
    /// mutation, so repeated reads between writes are O(1).
    #[must_use]
    pub fn total_count(&self) -> u64 {
        // Fast path: return the memoized cardinality if still valid. Read the
        // cache into a local and drop the guard before branching so the cache
        // lock is never held across the slow-path values read.
        let cached = *self
            .total_count_cache
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if let Some(count) = cached {
            return count;
        }
        // Slow path: recompute the union once and memoize it. Compute under the
        // values read lock so the counted snapshot is internally consistent.
        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);
        let count = union.len();
        *self
            .total_count_cache
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(count);
        count
    }

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

    /// Marker that introduces a bitmap-index key suffix.
    const SUFFIX_TAG: &'static [u8] = b"BMP:";

    /// Build the length-prefixed key suffix for `(field_name, value)`.
    ///
    /// Layout: `b"BMP:"` + `field_len: u16 BE` + `field_name` + `value`.
    ///
    /// # Errors
    ///
    /// Returns `IndexError::Serialization` if the field name exceeds
    /// `u16::MAX` bytes (which cannot be length-prefixed unambiguously).
    fn encode_suffix(field_name: &str, value: &str) -> Result<Vec<u8>, IndexError> {
        let field_len = u16::try_from(field_name.len()).map_err(|_| {
            IndexError::Serialization(format!(
                "bitmap index field name too long to length-prefix: {} bytes",
                field_name.len()
            ))
        })?;
        let mut suffix =
            Vec::with_capacity(Self::SUFFIX_TAG.len() + 2 + field_name.len() + value.len());
        suffix.extend_from_slice(Self::SUFFIX_TAG);
        suffix.extend_from_slice(&field_len.to_be_bytes());
        suffix.extend_from_slice(field_name.as_bytes());
        suffix.extend_from_slice(value.as_bytes());
        Ok(suffix)
    }

    /// Parse a key's suffix as `(field_name, value)` if it is a bitmap-index
    /// suffix for `expected_field`. Returns `Ok(None)` for non-matching keys.
    ///
    /// # Errors
    ///
    /// Returns `IndexError::Serialization` if a matching suffix is structurally
    /// corrupt (truncated length prefix or non-UTF-8 value bytes).
    fn parse_suffix_value<'a>(
        key: &'a [u8],
        expected_field: &str,
    ) -> Result<Option<&'a str>, IndexError> {
        // Skip the fixed encoded-key prefix (8 entity id + 1 NUL + 1 tag).
        if key.len() < ENCODED_KEY_PREFIX_LEN {
            return Ok(None);
        }
        let suffix = &key[ENCODED_KEY_PREFIX_LEN..];
        let Some(rest) = suffix.strip_prefix(Self::SUFFIX_TAG) else {
            return Ok(None);
        };
        if rest.len() < 2 {
            return Err(IndexError::Serialization(
                "truncated bitmap index field length prefix".to_string(),
            ));
        }
        let field_len = usize::from(u16::from_be_bytes([rest[0], rest[1]]));
        let body = &rest[2..];
        if body.len() < field_len {
            return Err(IndexError::Serialization(
                "bitmap index key shorter than declared field length".to_string(),
            ));
        }
        let (field_bytes, value_bytes) = body.split_at(field_len);
        if field_bytes != expected_field.as_bytes() {
            return Ok(None);
        }
        let value = std::str::from_utf8(value_bytes)
            .map_err(|e| IndexError::Serialization(e.to_string()))?;
        Ok(Some(value))
    }

    /// Serialize all bitmaps to (key, value) pairs for storage.
    ///
    /// Key: `encode_key(EntityId(INDEX_ROOT_ID), Tag::Idx, encode_suffix(field, value))`
    /// where `encode_suffix` length-prefixes the field name (see the type-level
    /// persistence docs for the exact byte layout).
    ///
    /// # Errors
    ///
    /// Returns `IndexError::Serialization` if bitmap serialization fails or the
    /// field name is too long to length-prefix.
    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 = Self::encode_suffix(&self.field_name, value)?;
            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(map);
        Ok(pairs)
    }

    /// Deserialize bitmaps from (key, value) pairs.
    ///
    /// Scans for length-prefixed `BMP:` keys whose embedded field name equals
    /// `field_name` and deserializes their bitmaps.
    ///
    /// # Errors
    ///
    /// Returns `IndexError::Serialization` if a matching key is structurally
    /// corrupt or a bitmap fails to deserialize.
    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 mut map = HashMap::new();

        for (key, value_bytes) in pairs {
            if let Some(field_value) = Self::parse_suffix_value(&key, &field_name)? {
                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)),
            // Loaded fresh; force a recompute on first read.
            total_count_cache: Arc::new(RwLock::new(None)),
        })
    }
}

// `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;
