//! Embedding lifecycle operations: insert, update, delete.
//!
//! These sit between the entity write API (`write_item()`) and the raw
//! [`VectorIndex`] trait. Each operation enforces L2 normalization, persists
//! the source-of-truth vector in the entity store, and updates the HNSW index.

use super::{
    super::{VectorError, VectorIndex},
    normalize::l2_normalize,
    serde::{embedding_store_key, serialize_embedding},
};
use crate::{
    schema::EntityId,
    storage::{StorageEngine, Tag, encode_key},
};

/// Suffix prefix marking a durable soft-delete (archive) tombstone for an
/// embedding slot, stored under [`Tag::Meta`] as `ARCHIVED:{slot_name}`.
///
/// This is the single source of truth for the tombstone marker shared between
/// the writer ([`delete_embedding`]) and the rebuild reader
/// (`EmbeddingSlotRegistry::rebuild_from_store`), so the two can never disagree
/// on the on-disk encoding.
pub(crate) const ARCHIVED_PREFIX: &[u8] = b"ARCHIVED:";

/// Build the entity-store key for an embedding slot's soft-delete tombstone.
///
/// Mirrors [`embedding_store_key`] but uses the [`ARCHIVED_PREFIX`] suffix so the
/// tombstone co-locates with the entity's other [`Tag::Meta`] keys while staying
/// distinct from the live `EMB:` embedding key.
#[must_use]
pub(crate) fn embedding_tombstone_key(entity_id: EntityId, slot_name: &str) -> Vec<u8> {
    let mut suffix = ARCHIVED_PREFIX.to_vec();
    suffix.extend_from_slice(slot_name.as_bytes());
    encode_key(entity_id, Tag::Meta, &suffix)
}

/// Validate, L2-normalize, and durably store an embedding in the entity store.
///
/// This is the shared write prologue for [`insert_embedding`] and
/// [`update_embedding`]: both must persist the source-of-truth f32 vector
/// before touching the (derived, rebuildable) HNSW index. Returns the
/// normalized vector so the caller can index it.
///
/// # Errors
///
/// - [`VectorError::DimensionMismatch`] if `raw_vector.len() != expected_dimensions`.
/// - [`VectorError::ZeroNormVector`] if the raw vector has zero norm.
/// - [`VectorError::Backend`] if the entity store write fails.
fn normalize_and_store(
    entity_id: EntityId,
    slot_name: &str,
    raw_vector: &[f32],
    expected_dimensions: usize,
    storage: &dyn StorageEngine,
) -> Result<Vec<f32>, VectorError> {
    if raw_vector.len() != expected_dimensions {
        return Err(VectorError::DimensionMismatch {
            expected: expected_dimensions,
            got: raw_vector.len(),
        });
    }

    let normalized = l2_normalize(raw_vector)?;

    // Store in entity store (source of truth) before indexing: the index is
    // derived state that can be rebuilt from this durable value on restart.
    let key = embedding_store_key(entity_id, slot_name);
    let value = serialize_embedding(&normalized);
    storage
        .put(&key, &value)
        .map_err(|e| VectorError::Backend(format!("entity store write failed: {e}")))?;

    // Clear any prior archive tombstone: writing an embedding un-archives the
    // slot. Without this, a slot that was soft-deleted then re-inserted would be
    // silently skipped by rebuild_from_store on the next restart due to the stale
    // tombstone. delete() is a no-op if no tombstone exists.
    let tombstone_key = embedding_tombstone_key(entity_id, slot_name);
    storage
        .delete(&tombstone_key)
        .map_err(|e| VectorError::Backend(format!("archive tombstone clear failed: {e}")))?;

    Ok(normalized)
}

/// Insert an embedding for an entity.
///
/// 1. Validates that `raw_vector.len() == expected_dimensions`.
/// 2. L2-normalizes the vector to unit length.
/// 3. Stores the normalized f32 vector in the entity store (source of truth).
/// 4. Inserts the normalized vector into the HNSW index.
///
/// # Errors
///
/// - [`VectorError::DimensionMismatch`] if dimensions do not match.
/// - [`VectorError::ZeroNormVector`] if the raw vector has zero norm.
/// - [`VectorError::Backend`] if the entity store write fails.
/// - Any error from [`VectorIndex::insert`].
pub fn insert_embedding(
    entity_id: EntityId,
    slot_name: &str,
    raw_vector: &[f32],
    expected_dimensions: usize,
    index: &dyn VectorIndex,
    storage: &dyn StorageEngine,
) -> Result<(), VectorError> {
    let normalized = normalize_and_store(
        entity_id,
        slot_name,
        raw_vector,
        expected_dimensions,
        storage,
    )?;

    // Insert into HNSW index
    index.insert(entity_id.as_u64(), &normalized)?;

    Ok(())
}

/// Update an embedding for an entity.
///
/// 1. Validates dimensions.
/// 2. L2-normalizes the new vector.
/// 3. Updates the entity store (overwrite) — durability-first.
/// 4. Tombstones the old vector in HNSW (ignores `NotFound` for first-time writes).
/// 5. Inserts the new vector into HNSW.
///
/// # Durability invariant (durable-first ordering)
///
/// Step 3 persists the new source-of-truth vector *before* either ANN mutation,
/// so the entity store and the derived HNSW index can only diverge in a way that
/// self-heals on restart. The two crash windows are:
///
/// - **After step 3, before/within steps 4–5:** the store already holds the new
///   vector while the ANN index still holds the old one (or, transiently between
///   4 and 5, neither). This window is microseconds and is acceptable per Spec
///   07, Section 6. On any crash here the in-memory ANN index is discarded and
///   `EmbeddingSlotRegistry::rebuild_from_store` re-derives it from the durable
///   store, which already reflects the new vector — the torn update is reconciled
///   with no manual repair.
/// - **If step 5 returns an error (ANN insert fails):** the store holds the new
///   vector but the live ANN index is missing the entity until the next restart's
///   rebuild. The error is propagated so the caller can retry; a retry re-runs
///   the idempotent store write and re-inserts into the index, converging the two.
///
/// The reverse order (touch ANN before persisting) would create a window where
/// the ANN index reflects an update the store never durably recorded, which a
/// rebuild would silently undo. Never index before persisting.
///
/// # Errors
///
/// - [`VectorError::DimensionMismatch`] if dimensions do not match.
/// - [`VectorError::ZeroNormVector`] if the raw vector has zero norm.
/// - [`VectorError::Backend`] if the entity store write fails.
/// - Any error from [`VectorIndex::insert`].
pub fn update_embedding(
    entity_id: EntityId,
    slot_name: &str,
    raw_vector: &[f32],
    expected_dimensions: usize,
    index: &dyn VectorIndex,
    storage: &dyn StorageEngine,
) -> Result<(), VectorError> {
    let normalized = normalize_and_store(
        entity_id,
        slot_name,
        raw_vector,
        expected_dimensions,
        storage,
    )?;

    // Tombstone old in HNSW, insert new.
    // delete() may return NotFound if the entity was never indexed (first embedding
    // via update path). That is fine -- ignore NotFound on the delete step.
    let delete_result = index.delete(entity_id.as_u64());
    match delete_result {
        Ok(()) | Err(VectorError::NotFound { .. }) => {}
        Err(e) => return Err(e),
    }
    index.insert(entity_id.as_u64(), &normalized)?;

    Ok(())
}

/// Delete an embedding for an entity.
///
/// 1. If `hard_delete` is true, removes the embedding from the entity store and
///    clears any archive tombstone (the slot is gone entirely).
/// 2. If `hard_delete` is false (archive / soft delete), keeps the source-of-truth
///    embedding but writes a durable `ARCHIVED:{slot}` tombstone so the entity does
///    not silently reappear in ANN results after a restart.
/// 3. Only after the durable mutation succeeds, tombstones the vector in HNSW.
///
/// # Durability invariant (durable-first ordering)
///
/// The in-memory ANN index is derived state rebuilt from the entity store on
/// restart (see `EmbeddingSlotRegistry::rebuild_from_store`). A soft delete that
/// only touched the ANN index would be undone by that rebuild — the archived
/// vector would resurface. The tombstone closes that divergence: rebuild skips
/// any `EMB:` key whose entity+slot has a matching tombstone.
///
/// The durable write happens *before* the ANN tombstone so the only crash window
/// is durable-deleted-but-still-ANN-resident. That state self-heals: the ANN
/// index is discarded on restart and rebuilt from the store, which already
/// reflects the delete (hard) or the archive tombstone (soft). The reverse order
/// would leave a crash window where the ANN tombstone is applied but the durable
/// tombstone is not — and rebuild would resurrect the vector. Never apply the ANN
/// tombstone first.
///
/// # Errors
///
/// - [`VectorError::NotFound`] if the vector is not in the HNSW index.
/// - [`VectorError::Backend`] if an entity store write or delete fails.
pub fn delete_embedding(
    entity_id: EntityId,
    slot_name: &str,
    index: &dyn VectorIndex,
    storage: &dyn StorageEngine,
    hard_delete: bool,
) -> Result<(), VectorError> {
    let tombstone_key = embedding_tombstone_key(entity_id, slot_name);

    // Durable mutation first: a crash after this point but before the ANN
    // tombstone is reconciled by rebuild_from_store, which reflects this write.
    if hard_delete {
        // Remove the source-of-truth embedding. Also clear any prior archive
        // tombstone so a future re-insert of this slot is not permanently
        // suppressed by a stale marker.
        let key = embedding_store_key(entity_id, slot_name);
        storage
            .delete(&key)
            .map_err(|e| VectorError::Backend(format!("entity store delete failed: {e}")))?;
        storage
            .delete(&tombstone_key)
            .map_err(|e| VectorError::Backend(format!("archive tombstone delete failed: {e}")))?;
    } else {
        // Soft delete / archive: keep the durable embedding but mark it archived
        // so rebuild_from_store skips it. The value is a zero-length sentinel —
        // only the key's presence matters.
        storage
            .put(&tombstone_key, &[])
            .map_err(|e| VectorError::Backend(format!("archive tombstone write failed: {e}")))?;
    }

    // Tombstone in HNSW (the derived index) only after the durable change lands.
    index.delete(entity_id.as_u64())?;

    Ok(())
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::cast_precision_loss)]
mod tests {
    use super::{
        super::serde::{deserialize_embedding, embedding_store_key},
        *,
    };
    use crate::{
        schema::EntityId,
        storage::{
            memory::InMemoryBackend,
            vector::{BruteForceIndex, VectorIndexConfig},
        },
    };

    // -------------------------------------------------------------------
    // insert_embedding
    // -------------------------------------------------------------------

    #[test]
    fn insert_embedding_validates_dimensions() {
        let config = VectorIndexConfig {
            dimensions: 3,
            ..VectorIndexConfig::default()
        };
        let index = BruteForceIndex::new(config);
        let storage = InMemoryBackend::new();

        let result = insert_embedding(
            EntityId::new(1),
            "content",
            &[1.0, 2.0],
            3,
            &index,
            &storage,
        );
        assert!(matches!(
            result,
            Err(VectorError::DimensionMismatch {
                expected: 3,
                got: 2
            })
        ));
    }

    #[test]
    fn insert_embedding_stores_normalized_vector() {
        let config = VectorIndexConfig {
            dimensions: 3,
            ..VectorIndexConfig::default()
        };
        let index = BruteForceIndex::new(config);
        let storage = InMemoryBackend::new();

        insert_embedding(
            EntityId::new(1),
            "content",
            &[3.0, 4.0, 0.0],
            3,
            &index,
            &storage,
        )
        .unwrap();

        // Read from entity store
        let key = embedding_store_key(EntityId::new(1), "content");
        let bytes = storage.get(&key).unwrap().unwrap();
        let stored = deserialize_embedding(&bytes).unwrap();

        // Should be normalized (norm = 5, so [0.6, 0.8, 0.0])
        assert!((stored[0] - 0.6).abs() < 1e-5);
        assert!((stored[1] - 0.8).abs() < 1e-5);
        assert!(stored[2].abs() < 1e-5);
    }

    #[test]
    fn insert_embedding_zero_vector_fails() {
        let config = VectorIndexConfig {
            dimensions: 3,
            ..VectorIndexConfig::default()
        };
        let index = BruteForceIndex::new(config);
        let storage = InMemoryBackend::new();

        let result = insert_embedding(
            EntityId::new(1),
            "content",
            &[0.0, 0.0, 0.0],
            3,
            &index,
            &storage,
        );
        assert!(matches!(result, Err(VectorError::ZeroNormVector)));
    }

    // -------------------------------------------------------------------
    // update_embedding
    // -------------------------------------------------------------------

    #[test]
    fn update_embedding_replaces_vector() {
        let config = VectorIndexConfig {
            dimensions: 3,
            ..VectorIndexConfig::default()
        };
        let index = BruteForceIndex::new(config);
        let storage = InMemoryBackend::new();

        // Insert original
        insert_embedding(
            EntityId::new(1),
            "content",
            &[1.0, 0.0, 0.0],
            3,
            &index,
            &storage,
        )
        .unwrap();

        // Update
        update_embedding(
            EntityId::new(1),
            "content",
            &[0.0, 1.0, 0.0],
            3,
            &index,
            &storage,
        )
        .unwrap();

        // Search should find the updated vector
        let results = index.search(&[0.0, 1.0, 0.0], 1, 0).unwrap();
        assert_eq!(results[0].id, 1);
        assert!(results[0].distance < 1e-5, "should match updated vector");
    }

    // -------------------------------------------------------------------
    // delete_embedding
    // -------------------------------------------------------------------

    #[test]
    fn delete_embedding_removes_from_index() {
        let config = VectorIndexConfig {
            dimensions: 3,
            ..VectorIndexConfig::default()
        };
        let index = BruteForceIndex::new(config);
        let storage = InMemoryBackend::new();

        insert_embedding(
            EntityId::new(1),
            "content",
            &[1.0, 0.0, 0.0],
            3,
            &index,
            &storage,
        )
        .unwrap();

        delete_embedding(EntityId::new(1), "content", &index, &storage, false).unwrap();

        // Should not appear in search results
        let results = index.search(&[1.0, 0.0, 0.0], 10, 0).unwrap();
        assert!(results.is_empty());

        // Soft delete: entity store still has the embedding (source of truth)...
        let key = embedding_store_key(EntityId::new(1), "content");
        assert!(storage.get(&key).unwrap().is_some());
        // ...but a durable archive tombstone is now present so a rebuild on the
        // next restart will skip it instead of resurrecting the soft-deleted vector.
        let tombstone = embedding_tombstone_key(EntityId::new(1), "content");
        assert!(
            storage.get(&tombstone).unwrap().is_some(),
            "soft delete must write a durable tombstone"
        );
    }

    #[test]
    fn hard_delete_clears_tombstone() {
        let config = VectorIndexConfig {
            dimensions: 3,
            ..VectorIndexConfig::default()
        };
        let index = BruteForceIndex::new(config);
        let storage = InMemoryBackend::new();

        insert_embedding(
            EntityId::new(1),
            "content",
            &[1.0, 0.0, 0.0],
            3,
            &index,
            &storage,
        )
        .unwrap();
        // Soft delete writes a tombstone.
        delete_embedding(EntityId::new(1), "content", &index, &storage, false).unwrap();
        let tombstone = embedding_tombstone_key(EntityId::new(1), "content");
        assert!(storage.get(&tombstone).unwrap().is_some());

        // Hard delete must remove BOTH the embedding and the stale tombstone, so
        // a future re-insert of this slot is not permanently suppressed.
        index.insert(1, &[1.0, 0.0, 0.0]).unwrap(); // re-add to index so delete() finds it
        delete_embedding(EntityId::new(1), "content", &index, &storage, true).unwrap();
        let key = embedding_store_key(EntityId::new(1), "content");
        assert!(storage.get(&key).unwrap().is_none(), "embedding removed");
        assert!(
            storage.get(&tombstone).unwrap().is_none(),
            "hard delete must clear the tombstone"
        );
    }

    #[test]
    fn reinsert_after_soft_delete_clears_tombstone() {
        let config = VectorIndexConfig {
            dimensions: 3,
            ..VectorIndexConfig::default()
        };
        let index = BruteForceIndex::new(config);
        let storage = InMemoryBackend::new();

        insert_embedding(
            EntityId::new(1),
            "content",
            &[1.0, 0.0, 0.0],
            3,
            &index,
            &storage,
        )
        .unwrap();
        delete_embedding(EntityId::new(1), "content", &index, &storage, false).unwrap();
        let tombstone = embedding_tombstone_key(EntityId::new(1), "content");
        assert!(storage.get(&tombstone).unwrap().is_some());

        // Re-inserting an embedding un-archives the slot: the tombstone must be
        // cleared, otherwise rebuild_from_store would silently skip the fresh
        // embedding on the next restart.
        insert_embedding(
            EntityId::new(1),
            "content",
            &[0.0, 1.0, 0.0],
            3,
            &index,
            &storage,
        )
        .unwrap();
        assert!(
            storage.get(&tombstone).unwrap().is_none(),
            "re-insert must clear the archive tombstone"
        );
    }

    #[test]
    fn soft_delete_durable_failure_leaves_ann_index_intact() {
        // W24 regression: the durable archive tombstone must be written BEFORE the
        // ANN tombstone. If the durable write fails, the ANN index must be left
        // untouched — otherwise a crash window would let rebuild_from_store
        // resurrect a vector the ANN had already deleted. We inject a backend
        // whose `put` always fails (the tombstone write) and assert the vector
        // is still present in the index after the error.
        use std::sync::Arc;

        use crate::storage::{StorageError, batch::WriteBatch, iterator::PrefixIterator};

        /// Backend that fails every `put` but delegates everything else, so the
        /// soft-delete tombstone write fails while reads/deletes still work.
        struct PutFailsBackend(Arc<InMemoryBackend>);

        impl StorageEngine for PutFailsBackend {
            fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
                self.0.get(key)
            }
            fn put(&self, _key: &[u8], _value: &[u8]) -> Result<(), StorageError> {
                Err(StorageError::Closed)
            }
            fn delete(&self, key: &[u8]) -> Result<(), StorageError> {
                self.0.delete(key)
            }
            fn scan_prefix(&self, prefix: &[u8]) -> PrefixIterator<'_> {
                self.0.scan_prefix(prefix)
            }
            fn write_batch(&self, batch: WriteBatch) -> Result<(), StorageError> {
                self.0.write_batch(batch)
            }
            fn flush(&self) -> Result<(), StorageError> {
                self.0.flush()
            }
        }

        let config = VectorIndexConfig {
            dimensions: 3,
            ..VectorIndexConfig::default()
        };
        let index = BruteForceIndex::new(config);

        // Seed the embedding through a working backend, then index it directly so
        // the ANN index holds the vector independent of the failing wrapper.
        let inner = Arc::new(InMemoryBackend::new());
        insert_embedding(
            EntityId::new(1),
            "content",
            &[1.0, 0.0, 0.0],
            3,
            &index,
            inner.as_ref(),
        )
        .unwrap();
        assert_eq!(index.len(), 1, "precondition: vector is indexed");

        // Soft delete against a backend whose tombstone `put` fails.
        let failing = PutFailsBackend(Arc::clone(&inner));
        let result = delete_embedding(EntityId::new(1), "content", &index, &failing, false);

        assert!(
            matches!(result, Err(VectorError::Backend(_))),
            "soft delete must surface the durable tombstone write failure"
        );
        // Durable-first ordering: because the tombstone write failed, the ANN
        // tombstone must NOT have been applied — the vector is still searchable.
        assert_eq!(
            index.len(),
            1,
            "ANN index must be untouched when the durable tombstone write fails"
        );
        let results = index.search(&[1.0, 0.0, 0.0], 1, 0).unwrap();
        assert_eq!(results[0].id, 1, "vector must still be in the ANN index");
    }

    #[test]
    fn delete_embedding_hard_removes_from_store() {
        let config = VectorIndexConfig {
            dimensions: 3,
            ..VectorIndexConfig::default()
        };
        let index = BruteForceIndex::new(config);
        let storage = InMemoryBackend::new();

        insert_embedding(
            EntityId::new(1),
            "content",
            &[1.0, 0.0, 0.0],
            3,
            &index,
            &storage,
        )
        .unwrap();

        delete_embedding(EntityId::new(1), "content", &index, &storage, true).unwrap();

        // Entity store should not have the embedding
        let key = embedding_store_key(EntityId::new(1), "content");
        assert!(storage.get(&key).unwrap().is_none());
    }

    // -------------------------------------------------------------------
    // Property tests
    // -------------------------------------------------------------------

    mod proptests {
        use proptest::prelude::*;

        use super::{
            super::super::serde::{deserialize_embedding, embedding_store_key},
            *,
        };

        proptest! {
            #[test]
            fn insert_embedding_searchable(
                dim in 2usize..64,
                n in 1usize..50,
            ) {
                let config = VectorIndexConfig {
                    dimensions: dim,
                    ..VectorIndexConfig::default()
                };
                let index = BruteForceIndex::new(config);
                let storage = InMemoryBackend::new();

                for id in 0..n as u64 {
                    let raw: Vec<f32> = (0..dim)
                        .map(|i| ((id as usize + i) % 100) as f32 / 100.0 + 0.01)
                        .collect();
                    insert_embedding(
                        EntityId::new(id + 1),
                        "content",
                        &raw,
                        dim,
                        &index,
                        &storage,
                    )
                    .unwrap();
                }

                // Verify all are searchable
                prop_assert_eq!(index.len(), n);

                // Verify entity store has the normalized vectors
                for id in 0..n as u64 {
                    let key = embedding_store_key(EntityId::new(id + 1), "content");
                    let bytes = storage.get(&key).unwrap();
                    prop_assert!(
                        bytes.is_some(),
                        "entity store should have embedding for id {id}"
                    );
                    let stored = deserialize_embedding(&bytes.unwrap()).unwrap();
                    let norm: f32 = stored.iter().map(|x| x * x).sum::<f32>().sqrt();
                    prop_assert!(
                        (1.0 - norm).abs() < 1e-5,
                        "stored embedding should be normalized"
                    );
                }
            }
        }
    }
}
