//! In-memory cache for memory entries, sessions, and knowledge graph.

use crate::vector_search;
use clawhdf5_format::float16::round_to_f16;

/// Every entry's embedding, in one contiguous `[N x dim]` buffer.
///
/// Rows are always exactly `dim` long: a shorter one is zero-padded, a longer
/// one truncated. The previous `Vec<Vec<f32>>` allowed ragged rows, which
/// silently misaligned the flattened copy that the batched kernels read — a
/// single wrong-length embedding shifted every row after it. Padding makes
/// that unrepresentable. A record stored without an embedding therefore holds
/// a zero row, and is told apart by its norm being zero rather than by length.
///
/// This used to be two fields — a `Vec<Vec<f32>>` and a flattened copy kept in
/// lock-step — which stored the whole corpus twice and cost one heap
/// allocation per entry on top. At 100k 384-dim entries that duplicate was
/// ~150 MiB. Indexing yields a `&[f32]` row, so `embeddings[i]` still reads
/// the same way.
#[derive(Debug, Clone, Default)]
pub struct Embeddings {
    flat: Vec<f32>,
    dim: usize,
}

impl Embeddings {
    pub fn new(dim: usize) -> Self {
        Self {
            flat: Vec::new(),
            dim,
        }
    }

    /// Number of embeddings.
    pub fn len(&self) -> usize {
        self.flat.len().checked_div(self.dim).unwrap_or(0)
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// The whole buffer, `[N x dim]` row-major — what batched kernels read.
    pub fn as_flat(&self) -> &[f32] {
        &self.flat
    }

    pub fn dim(&self) -> usize {
        self.dim
    }

    /// Row `i`, or `None` if out of range.
    pub fn get(&self, i: usize) -> Option<&[f32]> {
        let start = i.checked_mul(self.dim)?;
        self.flat.get(start..start.checked_add(self.dim)?)
    }

    pub fn iter(&self) -> impl ExactSizeIterator<Item = &[f32]> {
        self.flat.chunks_exact(self.dim.max(1))
    }

    /// Append one embedding. A row whose length doesn't match `dim` is padded
    /// or truncated, so the buffer stays rectangular whatever a caller passes.
    pub fn push(&mut self, embedding: &[f32]) {
        if self.dim == 0 {
            return;
        }
        let take = embedding.len().min(self.dim);
        self.flat.extend_from_slice(&embedding[..take]);
        self.flat.resize(self.flat.len() + (self.dim - take), 0.0);
    }

    /// Replace row `i`. Out-of-range indices are ignored.
    pub fn set(&mut self, i: usize, embedding: &[f32]) {
        let Some(start) = i.checked_mul(self.dim) else {
            return;
        };
        if start + self.dim > self.flat.len() {
            return;
        }
        let take = embedding.len().min(self.dim);
        self.flat[start..start + take].copy_from_slice(&embedding[..take]);
        self.flat[start + take..start + self.dim].fill(0.0);
    }

    /// Keep only the rows `keep` returns true for, preserving order.
    pub fn retain(&mut self, mut keep: impl FnMut(usize) -> bool) {
        if self.dim == 0 {
            return;
        }
        let mut write = 0usize;
        for read in 0..self.len() {
            if keep(read) {
                if write != read {
                    let (dst, src) = (write * self.dim, read * self.dim);
                    self.flat.copy_within(src..src + self.dim, dst);
                }
                write += 1;
            }
        }
        self.flat.truncate(write * self.dim);
    }

    /// Replace the contents with `rows`.
    pub fn reset_from(&mut self, dim: usize, rows: impl IntoIterator<Item = Vec<f32>>) {
        self.dim = dim;
        self.flat.clear();
        for row in rows {
            self.push(&row);
        }
    }

    /// Adopt an already-flat buffer, trimming any partial trailing row.
    pub fn set_flat(&mut self, dim: usize, mut flat: Vec<f32>) {
        self.dim = dim;
        match flat.len().checked_div(dim) {
            Some(rows) => flat.truncate(rows * dim),
            None => flat.clear(),
        }
        self.flat = flat;
    }
}

impl PartialEq for Embeddings {
    fn eq(&self, other: &Self) -> bool {
        self.dim == other.dim && self.flat == other.flat
    }
}

impl std::ops::Index<usize> for Embeddings {
    type Output = [f32];

    fn index(&self, i: usize) -> &[f32] {
        self.get(i).expect("embedding index out of range")
    }
}

/// In-memory cache for the /memory group data.
#[derive(Debug, Clone)]
pub struct MemoryCache {
    pub chunks: Vec<String>,
    pub embeddings: Embeddings,
    pub source_channels: Vec<String>,
    pub timestamps: Vec<f64>,
    pub session_ids: Vec<String>,
    pub tags: Vec<String>,
    pub tombstones: Vec<u8>,
    pub embedding_dim: usize,
    /// Pre-computed L2 norms for each embedding.
    pub norms: Vec<f32>,
    /// Hebbian activation weights (default 1.0 per entry).
    pub activation_weights: Vec<f32>,
    /// Round every embedding to IEEE half precision as it enters the cache,
    /// so the cache holds exactly what a `float16` store writes to disk. Set
    /// it with [`MemoryCache::set_half_precision`], which also rounds the
    /// rows already held.
    pub half_precision: bool,
}

impl MemoryCache {
    pub fn new(embedding_dim: usize) -> Self {
        Self {
            chunks: Vec::new(),
            embeddings: Embeddings::new(embedding_dim),
            source_channels: Vec::new(),
            timestamps: Vec::new(),
            session_ids: Vec::new(),
            tags: Vec::new(),
            tombstones: Vec::new(),
            embedding_dim,
            norms: Vec::new(),
            activation_weights: Vec::new(),
            half_precision: false,
        }
    }

    /// Switch half-precision rounding on or off. Turning it on rounds every
    /// embedding already held (and recomputes norms where one changed) —
    /// e.g. a `float16` store whose last checkpoint predates half-precision
    /// storage and so is still `f32` on disk.
    pub fn set_half_precision(&mut self, on: bool) {
        self.half_precision = on;
        if !on {
            return;
        }
        for i in 0..self.embeddings.len() {
            let row = &self.embeddings[i];
            if row
                .iter()
                .all(|&v| round_to_f16(v).to_bits() == v.to_bits())
            {
                continue;
            }
            let rounded: Vec<f32> = row.iter().map(|&v| round_to_f16(v)).collect();
            self.norms[i] = vector_search::compute_norm(&rounded);
            self.embeddings.set(i, &rounded);
        }
    }

    /// The embedding as the cache will hold it: rounded to half precision
    /// when [`Self::half_precision`] is on, otherwise unchanged.
    fn stored_form(&self, mut embedding: Vec<f32>) -> Vec<f32> {
        if self.half_precision {
            for v in &mut embedding {
                *v = round_to_f16(*v);
            }
        }
        embedding
    }

    /// Kept for callers that used to have to re-flatten after a bulk load.
    /// The buffer is always flat now, so there is nothing to rebuild.
    #[deprecated(note = "embeddings are stored flat; this is a no-op")]
    pub fn rebuild_flat(&mut self) {}

    /// The embeddings as one contiguous `[N x dim]` buffer.
    pub fn flat_embeddings(&self) -> &[f32] {
        self.embeddings.as_flat()
    }

    /// Total number of entries (including tombstoned).
    pub fn len(&self) -> usize {
        self.chunks.len()
    }

    pub fn is_empty(&self) -> bool {
        self.chunks.is_empty()
    }

    /// Number of active (non-tombstoned) entries.
    pub fn count_active(&self) -> usize {
        self.tombstones.iter().filter(|&&t| t == 0).count()
    }

    /// Push a new entry, returns its index.
    pub fn push(
        &mut self,
        chunk: String,
        embedding: Vec<f32>,
        source_channel: String,
        timestamp: f64,
        session_id: String,
        tags: String,
    ) -> usize {
        let idx = self.chunks.len();
        let embedding = self.stored_form(embedding);
        let norm = vector_search::compute_norm(&embedding);
        self.chunks.push(chunk);
        self.embeddings.push(&embedding);
        self.source_channels.push(source_channel);
        self.timestamps.push(timestamp);
        self.session_ids.push(session_id);
        self.tags.push(tags);
        self.tombstones.push(0);
        self.norms.push(norm);
        self.activation_weights.push(1.0);
        idx
    }

    /// Find an active (non-tombstoned) entry by tags (used as key for dedup).
    /// Returns the index of the first matching active entry, or None.
    pub fn find_by_tags(&self, tags: &str) -> Option<usize> {
        if tags.is_empty() {
            return None;
        }
        for (i, t) in self.tags.iter().enumerate() {
            if self.tombstones.get(i).copied().unwrap_or(1) == 0 && t == tags {
                return Some(i);
            }
        }
        None
    }

    /// Update an existing entry in-place (for upsert dedup).
    pub fn update(
        &mut self,
        idx: usize,
        chunk: String,
        embedding: Vec<f32>,
        source_channel: String,
        timestamp: f64,
        session_id: String,
    ) {
        if idx < self.chunks.len() {
            let embedding = self.stored_form(embedding);
            let norm = vector_search::compute_norm(&embedding);
            self.chunks[idx] = chunk;
            self.embeddings.set(idx, &embedding);
            self.source_channels[idx] = source_channel;
            self.timestamps[idx] = timestamp;
            self.session_ids[idx] = session_id;
            self.norms[idx] = norm;
            self.activation_weights[idx] = 1.0; // reset activation on update
        }
    }

    /// Mark an entry as deleted (tombstoned).
    pub fn mark_deleted(&mut self, id: usize) -> bool {
        if id < self.tombstones.len() && self.tombstones[id] == 0 {
            self.tombstones[id] = 1;
            true
        } else {
            false
        }
    }

    /// Fraction of entries that are tombstoned.
    pub fn tombstone_fraction(&self) -> f32 {
        if self.chunks.is_empty() {
            return 0.0;
        }
        let tombstoned = self.tombstones.iter().filter(|&&t| t == 1).count();
        tombstoned as f32 / self.chunks.len() as f32
    }

    /// Remove all tombstoned entries, returns number removed.
    /// Also returns a mapping from old indices to new indices (None if removed).
    /// Recomputes norms for remaining entries.
    pub fn compact(&mut self) -> (usize, Vec<Option<usize>>) {
        let old_len = self.chunks.len();
        let mut index_map = vec![None; old_len];
        let mut new_idx = 0usize;

        let mut new_chunks = Vec::new();
        let mut new_embeddings = Vec::new();
        let mut new_source_channels = Vec::new();
        let mut new_timestamps = Vec::new();
        let mut new_session_ids = Vec::new();
        let mut new_tags = Vec::new();
        let mut new_tombstones = Vec::new();
        let mut new_norms = Vec::new();
        let mut new_activation_weights = Vec::new();

        for (i, slot) in index_map.iter_mut().enumerate() {
            if self.tombstones[i] == 0 {
                *slot = Some(new_idx);
                new_idx += 1;
                let norm = vector_search::compute_norm(&self.embeddings[i]);
                new_chunks.push(self.chunks[i].clone());
                new_embeddings.push(self.embeddings[i].to_vec());
                new_source_channels.push(self.source_channels[i].clone());
                new_timestamps.push(self.timestamps[i]);
                new_session_ids.push(self.session_ids[i].clone());
                new_tags.push(self.tags[i].clone());
                new_tombstones.push(0u8);
                new_norms.push(norm);
                new_activation_weights.push(self.activation_weights[i]);
            }
        }

        let removed = old_len - new_chunks.len();
        self.chunks = new_chunks;
        self.embeddings
            .reset_from(self.embedding_dim, new_embeddings);
        self.source_channels = new_source_channels;
        self.timestamps = new_timestamps;
        self.session_ids = new_session_ids;
        self.tags = new_tags;
        self.tombstones = new_tombstones;
        self.norms = new_norms;
        self.activation_weights = new_activation_weights;

        (removed, index_map)
    }

    /// All embeddings as one owned `[N x dim]` buffer, for HDF5 storage.
    /// Prefer [`MemoryCache::flat_embeddings`] where a borrow will do.
    pub fn flat_embeddings_owned(&self) -> Vec<f32> {
        self.embeddings.as_flat().to_vec()
    }
}

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

    /// `embeddings_flat` must always equal a from-scratch flatten of `embeddings`.
    fn assert_flat_in_sync(cache: &MemoryCache) {
        let expected: Vec<f32> = cache.embeddings.iter().flatten().copied().collect();
        assert_eq!(cache.embeddings.as_flat(), expected);
    }

    #[test]
    fn push_keeps_flat_buffer_in_sync() {
        let mut cache = MemoryCache::new(3);
        cache.push(
            "a".into(),
            vec![1.0, 2.0, 3.0],
            "chan".into(),
            0.0,
            "s1".into(),
            String::new(),
        );
        cache.push(
            "b".into(),
            vec![4.0, 5.0, 6.0],
            "chan".into(),
            1.0,
            "s1".into(),
            String::new(),
        );
        assert_flat_in_sync(&cache);
        assert_eq!(
            cache.embeddings.as_flat(),
            vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
        );
    }

    #[test]
    fn update_keeps_flat_buffer_in_sync() {
        let mut cache = MemoryCache::new(3);
        cache.push(
            "a".into(),
            vec![1.0, 2.0, 3.0],
            "chan".into(),
            0.0,
            "s1".into(),
            String::new(),
        );
        cache.push(
            "b".into(),
            vec![4.0, 5.0, 6.0],
            "chan".into(),
            1.0,
            "s1".into(),
            String::new(),
        );
        cache.update(
            0,
            "a2".into(),
            vec![7.0, 8.0, 9.0],
            "chan".into(),
            2.0,
            "s1".into(),
        );
        assert_flat_in_sync(&cache);
        assert_eq!(
            cache.embeddings.as_flat(),
            vec![7.0, 8.0, 9.0, 4.0, 5.0, 6.0],
            "update must overwrite the correct flat slice, not just append"
        );
    }

    #[test]
    fn compact_keeps_flat_buffer_in_sync() {
        let mut cache = MemoryCache::new(2);
        cache.push(
            "a".into(),
            vec![1.0, 1.0],
            "chan".into(),
            0.0,
            "s1".into(),
            String::new(),
        );
        cache.push(
            "b".into(),
            vec![2.0, 2.0],
            "chan".into(),
            1.0,
            "s1".into(),
            String::new(),
        );
        cache.push(
            "c".into(),
            vec![3.0, 3.0],
            "chan".into(),
            2.0,
            "s1".into(),
            String::new(),
        );
        cache.mark_deleted(1);
        cache.compact();
        assert_flat_in_sync(&cache);
        assert_eq!(cache.embeddings.as_flat(), vec![1.0, 1.0, 3.0, 3.0]);
    }

    #[test]
    fn rebuild_flat_matches_manual_flatten() {
        let mut cache = MemoryCache::new(2);
        cache
            .embeddings
            .reset_from(2, vec![vec![1.0, 2.0], vec![3.0, 4.0]]);
        assert_eq!(cache.embeddings.as_flat(), vec![1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    fn set_half_precision_rounds_existing_rows_and_their_norms() {
        // A store with float16 set whose checkpoint is still f32 on disk
        // loads full-precision rows; switching rounding on must bring them to
        // exactly what the next checkpoint will write.
        let mut cache = MemoryCache::new(3);
        cache.push(
            "a".into(),
            vec![0.1, 0.2, 0.3],
            "c".into(),
            0.0,
            "s".into(),
            "".into(),
        );
        cache.push(
            "b".into(),
            vec![0.5, 0.25, 1.0],
            "c".into(),
            0.0,
            "s".into(),
            "".into(),
        );
        let exact_norm = cache.norms[0];

        cache.set_half_precision(true);
        let row0: Vec<f32> = [0.1f32, 0.2, 0.3]
            .iter()
            .map(|&v| round_to_f16(v))
            .collect();
        assert_eq!(&cache.embeddings[0], row0.as_slice());
        assert_eq!(cache.norms[0], vector_search::compute_norm(&row0));
        assert_ne!(cache.norms[0], exact_norm);
        // Already representable: untouched.
        assert_eq!(&cache.embeddings[1], &[0.5, 0.25, 1.0]);

        // New rows are rounded as they arrive, and updates too.
        cache.push(
            "c".into(),
            vec![0.1, 0.0, 0.0],
            "c".into(),
            0.0,
            "s".into(),
            "".into(),
        );
        assert_eq!(cache.embeddings[2][0], round_to_f16(0.1));
        cache.update(
            2,
            "c".into(),
            vec![0.3, 0.0, 0.0],
            "c".into(),
            0.0,
            "s".into(),
        );
        assert_eq!(cache.embeddings[2][0], round_to_f16(0.3));
    }
}