storage/

object.rs

//! Object Storage Backend using OpenDAL
//!
//! Provides storage abstraction for multiple cloud storage backends:
//! - S3-compatible (AWS S3, MinIO, DigitalOcean Spaces, etc.)
//! - Azure Blob Storage
//! - Google Cloud Storage (GCS)
//! - Local filesystem
//! - In-memory (for testing)

use async_trait::async_trait;
use common::{DakeraError, NamespaceId, Result, Vector, VectorId};
use futures_util::stream::{self, StreamExt};
use opendal::{layers::RetryLayer, services, Operator};
use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;

static TMP_COUNTER: AtomicU64 = AtomicU64::new(0);

use crate::traits::{IndexStorage, IndexType, VectorStorage};

fn s3_concurrent_ops() -> usize {
    std::env::var("DAKERA_S3_CONCURRENT_OPS")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(16)
}

/// Serializable namespace metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
struct NamespaceMetadata {
    dimension: Option<usize>,
    vector_count: usize,
    created_at: u64,
    updated_at: u64,
}

/// Serializable vector data for storage
#[derive(Debug, Clone, Serialize, Deserialize)]
struct StoredVector {
    id: String,
    values: Vec<f32>,
    metadata: Option<serde_json::Value>,
}

impl From<Vector> for StoredVector {
    fn from(v: Vector) -> Self {
        Self {
            id: v.id,
            values: v.values,
            metadata: v.metadata,
        }
    }
}

impl From<StoredVector> for Vector {
    fn from(v: StoredVector) -> Self {
        Self {
            id: v.id,
            values: v.values,
            metadata: v.metadata,
            ttl_seconds: None,
            expires_at: None,
        }
    }
}

/// Object storage backend configuration
#[derive(Debug, Clone, Default)]
pub enum ObjectStorageConfig {
    /// In-memory storage (for testing)
    #[default]
    Memory,
    /// Local filesystem storage
    Filesystem { root: String },
    /// S3-compatible storage (S3, MinIO, etc.)
    S3 {
        bucket: String,
        region: Option<String>,
        endpoint: Option<String>,
        access_key_id: Option<String>,
        secret_access_key: Option<String>,
    },
    /// Azure Blob Storage
    Azure {
        container: String,
        account_name: String,
        account_key: Option<String>,
        sas_token: Option<String>,
        endpoint: Option<String>,
    },
    /// Google Cloud Storage
    Gcs {
        bucket: String,
        credential_path: Option<String>,
        endpoint: Option<String>,
    },
}

/// Object storage backend using OpenDAL
#[derive(Clone)]
pub struct ObjectStorage {
    operator: Operator,
}

impl ObjectStorage {
    /// Create a new object storage with the given configuration
    pub fn new(config: ObjectStorageConfig) -> Result<Self> {
        let operator = Self::build_operator(&config)?;
        Ok(Self { operator })
    }

    /// Create in-memory storage (for testing)
    pub fn memory() -> Result<Self> {
        Self::new(ObjectStorageConfig::Memory)
    }

    /// Create filesystem storage
    pub fn filesystem(root: impl Into<String>) -> Result<Self> {
        Self::new(ObjectStorageConfig::Filesystem { root: root.into() })
    }

    /// Create S3 storage
    pub fn s3(bucket: impl Into<String>) -> Result<Self> {
        Self::new(ObjectStorageConfig::S3 {
            bucket: bucket.into(),
            region: None,
            endpoint: None,
            access_key_id: None,
            secret_access_key: None,
        })
    }

    /// Create Azure Blob storage
    pub fn azure(container: impl Into<String>, account_name: impl Into<String>) -> Result<Self> {
        Self::new(ObjectStorageConfig::Azure {
            container: container.into(),
            account_name: account_name.into(),
            account_key: None,
            sas_token: None,
            endpoint: None,
        })
    }

    /// Create Google Cloud Storage
    pub fn gcs(bucket: impl Into<String>) -> Result<Self> {
        Self::new(ObjectStorageConfig::Gcs {
            bucket: bucket.into(),
            credential_path: None,
            endpoint: None,
        })
    }

    pub fn build_operator(config: &ObjectStorageConfig) -> Result<Operator> {
        match config {
            ObjectStorageConfig::Memory => {
                let builder = services::Memory::default();
                Operator::new(builder)
                    .map(|op| op.finish())
                    .map_err(|e| DakeraError::Storage(e.to_string()))
            }
            ObjectStorageConfig::Filesystem { root } => {
                let builder = services::Fs::default().root(root);
                Operator::new(builder)
                    .map(|op| op.layer(Self::retry_layer()).finish())
                    .map_err(|e| DakeraError::Storage(e.to_string()))
            }
            ObjectStorageConfig::S3 {
                bucket,
                region,
                endpoint,
                access_key_id,
                secret_access_key,
            } => {
                let mut builder = services::S3::default().bucket(bucket);

                if let Some(region) = region {
                    builder = builder.region(region);
                }
                if let Some(endpoint) = endpoint {
                    builder = builder.endpoint(endpoint);
                }
                if let Some(key) = access_key_id {
                    builder = builder.access_key_id(key);
                }
                if let Some(secret) = secret_access_key {
                    builder = builder.secret_access_key(secret);
                }

                Operator::new(builder)
                    .map(|op| op.layer(Self::retry_layer()).finish())
                    .map_err(|e| DakeraError::Storage(e.to_string()))
            }
            ObjectStorageConfig::Azure {
                container,
                account_name,
                account_key,
                sas_token,
                endpoint,
            } => {
                let mut builder = services::Azblob::default()
                    .container(container)
                    .account_name(account_name);

                if let Some(key) = account_key {
                    builder = builder.account_key(key);
                }
                if let Some(token) = sas_token {
                    builder = builder.sas_token(token);
                }
                if let Some(endpoint) = endpoint {
                    builder = builder.endpoint(endpoint);
                }

                Operator::new(builder)
                    .map(|op| op.layer(Self::retry_layer()).finish())
                    .map_err(|e| DakeraError::Storage(e.to_string()))
            }
            ObjectStorageConfig::Gcs {
                bucket,
                credential_path,
                endpoint,
            } => {
                let mut builder = services::Gcs::default().bucket(bucket);

                if let Some(cred_path) = credential_path {
                    builder = builder.credential_path(cred_path);
                }
                if let Some(endpoint) = endpoint {
                    builder = builder.endpoint(endpoint);
                }

                Operator::new(builder)
                    .map(|op| op.layer(Self::retry_layer()).finish())
                    .map_err(|e| DakeraError::Storage(e.to_string()))
            }
        }
    }

    fn retry_layer() -> RetryLayer {
        // DAK-3430: default capped from 10 → 3 so a MinIO 429 storm fails fast
        // (≤3×60s ≈ 3min) rather than blocking production for 10min. Override
        // with DAKERA_S3_MAX_RETRIES=10 for high-resilience deployments.
        let max_times: usize = std::env::var("DAKERA_S3_MAX_RETRIES")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(3);
        let min_delay_ms: u64 = std::env::var("DAKERA_S3_RETRY_MIN_DELAY_MS")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(500);
        let max_delay_secs: u64 = std::env::var("DAKERA_S3_RETRY_MAX_DELAY_SECS")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(60);

        tracing::info!(
            max_times,
            min_delay_ms,
            max_delay_secs,
            "S3 retry layer configured"
        );

        RetryLayer::new()
            .with_max_times(max_times)
            .with_min_delay(Duration::from_millis(min_delay_ms))
            .with_max_delay(Duration::from_secs(max_delay_secs))
            .with_jitter()
            .with_factor(2.0)
    }

    /// Get the path for a vector
    fn vector_path(namespace: &str, vector_id: &str) -> String {
        format!("namespaces/{}/vectors/{}.json", namespace, vector_id)
    }

    /// Get the path for namespace metadata
    fn namespace_meta_path(namespace: &str) -> String {
        format!("namespaces/{}/meta.json", namespace)
    }

    /// Get the path prefix for all vectors in a namespace
    fn namespace_vectors_prefix(namespace: &str) -> String {
        format!("namespaces/{}/vectors/", namespace)
    }

    /// Read namespace metadata
    async fn read_namespace_meta(&self, namespace: &str) -> Result<Option<NamespaceMetadata>> {
        let path = Self::namespace_meta_path(namespace);
        match self.operator.read(&path).await {
            Ok(data) => {
                let bytes = data.to_vec();
                if bytes.is_empty() {
                    tracing::warn!(
                        namespace = %namespace,
                        path = %path,
                        "Empty namespace metadata file detected, treating as missing"
                    );
                    return Ok(None);
                }
                match serde_json::from_slice(&bytes) {
                    Ok(meta) => Ok(Some(meta)),
                    Err(e) => {
                        tracing::warn!(
                            namespace = %namespace,
                            path = %path,
                            error = %e,
                            bytes_len = bytes.len(),
                            "Corrupted namespace metadata, treating as missing and will be recreated"
                        );
                        Ok(None)
                    }
                }
            }
            Err(e) if e.kind() == opendal::ErrorKind::NotFound => Ok(None),
            Err(e) => Err(DakeraError::Storage(e.to_string())),
        }
    }

    /// Write namespace metadata atomically.
    ///
    /// Writes to a `.tmp` file first, then renames to the final path. On POSIX filesystems
    async fn write_namespace_meta(&self, namespace: &str, meta: &NamespaceMetadata) -> Result<()> {
        let path = Self::namespace_meta_path(namespace);
        let data = serde_json::to_vec(meta).map_err(|e| DakeraError::Storage(e.to_string()))?;
        self.write_atomic(&path, data).await
    }

    /// Write data to `path` atomically via tmp+rename (DAK-4545).
    ///
    /// Writes to a `.tmp` file first, then renames to the final path. On POSIX filesystems
    /// rename(2) is atomic within the same device, preventing concurrent readers from seeing
    /// a truncated/empty file during the write window (O_TRUNC race). For object-store
    /// backends that don't support rename (e.g. S3 without copy-delete), falls back to a
    /// direct write — S3 PUT is itself atomic (readers get previous version until PUT completes).
    async fn write_atomic(&self, path: &str, data: Vec<u8>) -> Result<()> {
        let seq = TMP_COUNTER.fetch_add(1, Ordering::Relaxed);
        let tmp_path = format!("{}.tmp.{}.{}", path, Self::now(), seq);
        let rename_result = async {
            self.operator
                .write(&tmp_path, data.clone())
                .await
                .map_err(|e| DakeraError::Storage(e.to_string()))?;
            self.operator
                .rename(&tmp_path, path)
                .await
                .map_err(|e| DakeraError::Storage(e.to_string()))?;
            Ok::<(), DakeraError>(())
        }
        .await;
        if let Err(e) = rename_result {
            // Backend doesn't support rename: fall back to direct write.
            tracing::debug!(path = %path, error = %e, "atomic rename failed, falling back to direct write");
            let _ = self.operator.delete(&tmp_path).await;
            self.operator
                .write(path, data)
                .await
                .map_err(|e| DakeraError::Storage(e.to_string()))?;
        }
        Ok(())
    }

    /// Get current timestamp
    fn now() -> u64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs()
    }
}

#[async_trait]
impl VectorStorage for ObjectStorage {
    async fn upsert(&self, namespace: &NamespaceId, vectors: Vec<Vector>) -> Result<usize> {
        if vectors.is_empty() {
            return Ok(0);
        }

        // Get or create namespace metadata
        let mut meta = self
            .read_namespace_meta(namespace)
            .await?
            .unwrap_or_else(|| NamespaceMetadata {
                dimension: None,
                vector_count: 0,
                created_at: Self::now(),
                updated_at: Self::now(),
            });

        // Validate dimensions
        let first_dim = vectors[0].values.len();
        if let Some(dim) = meta.dimension {
            for v in &vectors {
                if v.values.len() != dim {
                    return Err(DakeraError::DimensionMismatch {
                        expected: dim,
                        actual: v.values.len(),
                    });
                }
            }
        } else {
            meta.dimension = Some(first_dim);
        }

        let mut upserted = 0;
        for vector in vectors {
            let path = Self::vector_path(namespace, &vector.id);
            let stored: StoredVector = vector.into();
            let data =
                serde_json::to_vec(&stored).map_err(|e| DakeraError::Storage(e.to_string()))?;

            // Check if this is an insert or update
            let exists = self
                .operator
                .exists(&path)
                .await
                .map_err(|e| DakeraError::Storage(e.to_string()))?;

            self.write_atomic(&path, data).await?;

            if !exists {
                meta.vector_count += 1;
            }
            upserted += 1;
        }

        meta.updated_at = Self::now();
        self.write_namespace_meta(namespace, &meta).await?;

        tracing::debug!(
            namespace = namespace,
            upserted = upserted,
            "Upserted vectors to object storage"
        );

        Ok(upserted)
    }

    async fn get(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<Vec<Vector>> {
        if ids.is_empty() {
            return Ok(Vec::new());
        }

        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        let read_tasks: Vec<_> = ids
            .iter()
            .map(|id| {
                let operator = self.operator.clone();
                let path = Self::vector_path(namespace, id);
                let id = id.clone();
                async move {
                    match operator.read(&path).await {
                        Ok(data) => {
                            let bytes = data.to_vec();
                            if bytes.is_empty() {
                                tracing::warn!(
                                    vector_id = %id,
                                    "Empty vector file detected, skipping"
                                );
                                return Ok(None);
                            }
                            match serde_json::from_slice::<StoredVector>(&bytes) {
                                Ok(stored) => {
                                    let vector: Vector = stored.into();
                                    if !vector.is_expired_at(now) {
                                        Ok(Some(vector))
                                    } else {
                                        Ok(None)
                                    }
                                }
                                Err(e) => {
                                    tracing::warn!(
                                        vector_id = %id,
                                        error = %e,
                                        bytes_len = bytes.len(),
                                        "Corrupted vector file detected, skipping"
                                    );
                                    Ok(None)
                                }
                            }
                        }
                        Err(e) if e.kind() == opendal::ErrorKind::NotFound => Ok(None),
                        Err(e) => Err(DakeraError::Storage(e.to_string())),
                    }
                }
            })
            .collect();

        let results: Vec<Result<Option<Vector>>> = stream::iter(read_tasks)
            .buffer_unordered(s3_concurrent_ops())
            .collect()
            .await;

        let mut vectors = Vec::with_capacity(ids.len());
        for result in results {
            if let Some(v) = result? {
                vectors.push(v);
            }
        }
        Ok(vectors)
    }

    async fn get_all(&self, namespace: &NamespaceId) -> Result<Vec<Vector>> {
        let prefix = Self::namespace_vectors_prefix(namespace);

        let entries = self
            .operator
            .list(&prefix)
            .await
            .map_err(|e| DakeraError::Storage(e.to_string()))?;

        let json_paths: Vec<String> = entries
            .into_iter()
            .filter(|e| e.path().ends_with(".json"))
            .map(|e| e.path().to_string())
            .collect();

        if json_paths.is_empty() {
            return Ok(Vec::new());
        }

        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        let results: Vec<Option<Vector>> = stream::iter(json_paths.into_iter().map(|path| {
            let operator = self.operator.clone();
            async move {
                match operator.read(&path).await {
                    Ok(data) => {
                        let bytes = data.to_vec();
                        if let Ok(stored) = serde_json::from_slice::<StoredVector>(&bytes) {
                            let vector: Vector = stored.into();
                            if !vector.is_expired_at(now) {
                                return Some(vector);
                            }
                        }
                        None
                    }
                    Err(e) => {
                        tracing::warn!(path = %path, error = %e, "Failed to read vector");
                        None
                    }
                }
            }
        }))
        .buffer_unordered(s3_concurrent_ops())
        .collect()
        .await;

        Ok(results.into_iter().flatten().collect())
    }

    async fn delete(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<usize> {
        if ids.is_empty() {
            return Ok(0);
        }

        let delete_tasks: Vec<_> = ids
            .iter()
            .map(|id| {
                let operator = self.operator.clone();
                let path = Self::vector_path(namespace, id);
                async move {
                    let exists = operator
                        .exists(&path)
                        .await
                        .map_err(|e| DakeraError::Storage(e.to_string()))?;
                    if exists {
                        match operator.delete(&path).await {
                            Ok(_) => Ok(true),
                            Err(e) if e.kind() == opendal::ErrorKind::NotFound => Ok(false),
                            Err(e) => Err(DakeraError::Storage(e.to_string())),
                        }
                    } else {
                        Ok(false)
                    }
                }
            })
            .collect();

        let results: Vec<Result<bool>> = stream::iter(delete_tasks)
            .buffer_unordered(s3_concurrent_ops())
            .collect()
            .await;

        let mut deleted = 0;
        for result in results {
            if result? {
                deleted += 1;
            }
        }

        // Update metadata
        if deleted > 0 {
            if let Some(mut meta) = self.read_namespace_meta(namespace).await? {
                meta.vector_count = meta.vector_count.saturating_sub(deleted);
                meta.updated_at = Self::now();
                self.write_namespace_meta(namespace, &meta).await?;
            }
        }

        Ok(deleted)
    }

    async fn namespace_exists(&self, namespace: &NamespaceId) -> Result<bool> {
        Ok(self.read_namespace_meta(namespace).await?.is_some())
    }

    async fn ensure_namespace(&self, namespace: &NamespaceId) -> Result<()> {
        if self.read_namespace_meta(namespace).await?.is_none() {
            let meta = NamespaceMetadata {
                dimension: None,
                vector_count: 0,
                created_at: Self::now(),
                updated_at: Self::now(),
            };
            self.write_namespace_meta(namespace, &meta).await?;
        }
        Ok(())
    }

    async fn dimension(&self, namespace: &NamespaceId) -> Result<Option<usize>> {
        Ok(self
            .read_namespace_meta(namespace)
            .await?
            .and_then(|m| m.dimension))
    }

    async fn count(&self, namespace: &NamespaceId) -> Result<usize> {
        Ok(self
            .read_namespace_meta(namespace)
            .await?
            .map(|m| m.vector_count)
            .unwrap_or(0))
    }

    async fn list_namespaces(&self) -> Result<Vec<NamespaceId>> {
        let entries = self
            .operator
            .list("namespaces/")
            .await
            .map_err(|e| DakeraError::Storage(e.to_string()))?;

        let mut namespaces = Vec::new();
        for entry in entries {
            let path = entry.path();
            // Extract namespace name from path like "namespaces/myns/"
            if let Some(ns) = path.strip_prefix("namespaces/") {
                let ns = ns.trim_end_matches('/');
                if !ns.is_empty() && !ns.contains('/') {
                    // Only include namespaces that actually have metadata (meta.json).
                    // This filters out ghost directory entries left behind after
                    // namespace deletion on backends where empty directories persist
                    // (e.g., local filesystem, some S3-compatible stores).
                    if self.read_namespace_meta(ns).await?.is_some() {
                        namespaces.push(ns.to_string());
                    }
                }
            }
        }

        Ok(namespaces)
    }

    async fn delete_namespace(&self, namespace: &NamespaceId) -> Result<bool> {
        // Check if namespace exists
        if !self.namespace_exists(namespace).await? {
            return Ok(false);
        }

        // Recursively remove everything under the namespace prefix.
        // This deletes vectors, metadata, indexes, AND directory entries,
        // preventing ghost namespaces from appearing in list_namespaces().
        let prefix = format!("namespaces/{}/", namespace);
        self.operator
            .delete_with(&prefix)
            .recursive(true)
            .await
            .map_err(|e| DakeraError::Storage(e.to_string()))?;

        tracing::debug!(
            namespace = namespace,
            "Deleted namespace from object storage"
        );

        Ok(true)
    }

    async fn cleanup_expired(&self, _namespace: &NamespaceId) -> Result<usize> {
        // Object storage doesn't track TTL internally - cleanup handled by application layer
        // StoredVector doesn't persist TTL fields, so nothing to clean up here
        Ok(0)
    }

    async fn cleanup_all_expired(&self) -> Result<usize> {
        // Object storage doesn't track TTL internally - cleanup handled by application layer
        Ok(0)
    }
}

#[async_trait]
impl IndexStorage for ObjectStorage {
    async fn save_index(
        &self,
        namespace: &NamespaceId,
        index_type: IndexType,
        data: Vec<u8>,
    ) -> Result<()> {
        let path = format!(
            "namespaces/{}/indexes/{}.bin",
            namespace,
            index_type.as_str()
        );
        self.operator
            .write(&path, data)
            .await
            .map_err(|e| DakeraError::Storage(e.to_string()))?;

        tracing::debug!(
            namespace = namespace,
            index_type = index_type.as_str(),
            "Saved index to object storage"
        );
        Ok(())
    }

    async fn load_index(
        &self,
        namespace: &NamespaceId,
        index_type: IndexType,
    ) -> Result<Option<Vec<u8>>> {
        let path = format!(
            "namespaces/{}/indexes/{}.bin",
            namespace,
            index_type.as_str()
        );
        match self.operator.read(&path).await {
            Ok(data) => {
                tracing::debug!(
                    namespace = namespace,
                    index_type = index_type.as_str(),
                    size = data.len(),
                    "Loaded index from object storage"
                );
                Ok(Some(data.to_vec()))
            }
            Err(e) if e.kind() == opendal::ErrorKind::NotFound => Ok(None),
            Err(e) => Err(DakeraError::Storage(e.to_string())),
        }
    }

    async fn delete_index(&self, namespace: &NamespaceId, index_type: IndexType) -> Result<bool> {
        let path = format!(
            "namespaces/{}/indexes/{}.bin",
            namespace,
            index_type.as_str()
        );
        let exists = self
            .operator
            .exists(&path)
            .await
            .map_err(|e| DakeraError::Storage(e.to_string()))?;

        if exists {
            self.operator
                .delete(&path)
                .await
                .map_err(|e| DakeraError::Storage(e.to_string()))?;
            tracing::debug!(
                namespace = namespace,
                index_type = index_type.as_str(),
                "Deleted index from object storage"
            );
        }
        Ok(exists)
    }

    async fn index_exists(&self, namespace: &NamespaceId, index_type: IndexType) -> Result<bool> {
        let path = format!(
            "namespaces/{}/indexes/{}.bin",
            namespace,
            index_type.as_str()
        );
        self.operator
            .exists(&path)
            .await
            .map_err(|e| DakeraError::Storage(e.to_string()))
    }

    async fn list_indexes(&self, namespace: &NamespaceId) -> Result<Vec<IndexType>> {
        let prefix = format!("namespaces/{}/indexes/", namespace);
        let entries = self
            .operator
            .list(&prefix)
            .await
            .map_err(|e| DakeraError::Storage(e.to_string()))?;

        let mut indexes = Vec::new();
        for entry in entries {
            let path = entry.path();
            if path.ends_with(".bin") {
                // Extract index type from filename
                if let Some(filename) = path.strip_prefix(&prefix) {
                    let name = filename.trim_end_matches(".bin");
                    match name {
                        "hnsw" => indexes.push(IndexType::Hnsw),
                        "pq" => indexes.push(IndexType::Pq),
                        "ivf" => indexes.push(IndexType::Ivf),
                        "spfresh" => indexes.push(IndexType::SpFresh),
                        "fulltext" => indexes.push(IndexType::FullText),
                        _ => {} // Ignore unknown index types
                    }
                }
            }
        }

        Ok(indexes)
    }
}

/// Create an OpenDAL operator from configuration without constructing a full ObjectStorage.
/// Useful for lightweight S3 access (e.g., BackupManager metadata persistence).
pub fn create_operator(config: &ObjectStorageConfig) -> Result<Operator> {
    ObjectStorage::build_operator(config)
}

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

    #[tokio::test]
    async fn test_object_storage_memory() {
        let storage = ObjectStorage::memory().unwrap();
        let namespace = "test".to_string();

        // Ensure namespace
        storage.ensure_namespace(&namespace).await.unwrap();
        assert!(storage.namespace_exists(&namespace).await.unwrap());

        // Insert vectors
        let vectors = vec![
            Vector {
                id: "v1".to_string(),
                values: vec![1.0, 2.0, 3.0],
                metadata: None,
                ttl_seconds: None,
                expires_at: None,
            },
            Vector {
                id: "v2".to_string(),
                values: vec![4.0, 5.0, 6.0],
                metadata: Some(serde_json::json!({"key": "value"})),
                ttl_seconds: None,
                expires_at: None,
            },
        ];

        let count = storage.upsert(&namespace, vectors).await.unwrap();
        assert_eq!(count, 2);

        // Get single vector
        let results = storage.get(&namespace, &["v1".to_string()]).await.unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, "v1");
        assert_eq!(results[0].values, vec![1.0, 2.0, 3.0]);

        // Get all vectors
        let all = storage.get_all(&namespace).await.unwrap();
        assert_eq!(all.len(), 2);

        // Count
        assert_eq!(storage.count(&namespace).await.unwrap(), 2);

        // Delete
        let deleted = storage
            .delete(&namespace, &["v1".to_string()])
            .await
            .unwrap();
        assert_eq!(deleted, 1);
        assert!(storage
            .get(&namespace, &["v1".to_string()])
            .await
            .unwrap()
            .is_empty());
        assert_eq!(storage.count(&namespace).await.unwrap(), 1);
    }

    #[tokio::test]
    async fn test_object_storage_dimension_mismatch() {
        let storage = ObjectStorage::memory().unwrap();
        let namespace = "test".to_string();
        storage.ensure_namespace(&namespace).await.unwrap();

        // Insert first vector
        let v1 = vec![Vector {
            id: "v1".to_string(),
            values: vec![1.0, 2.0, 3.0],
            metadata: None,
            ttl_seconds: None,
            expires_at: None,
        }];
        storage.upsert(&namespace, v1).await.unwrap();

        // Try to insert vector with different dimension
        let v2 = vec![Vector {
            id: "v2".to_string(),
            values: vec![1.0, 2.0], // Wrong dimension
            metadata: None,
            ttl_seconds: None,
            expires_at: None,
        }];
        let result = storage.upsert(&namespace, v2).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_object_storage_upsert() {
        let storage = ObjectStorage::memory().unwrap();
        let namespace = "test".to_string();
        storage.ensure_namespace(&namespace).await.unwrap();

        // Insert
        let v1 = vec![Vector {
            id: "v1".to_string(),
            values: vec![1.0, 2.0],
            metadata: None,
            ttl_seconds: None,
            expires_at: None,
        }];
        storage.upsert(&namespace, v1).await.unwrap();

        // Upsert (update)
        let v1_updated = vec![Vector {
            id: "v1".to_string(),
            values: vec![3.0, 4.0],
            metadata: None,
            ttl_seconds: None,
            expires_at: None,
        }];
        storage.upsert(&namespace, v1_updated).await.unwrap();

        // Verify update
        let results = storage.get(&namespace, &["v1".to_string()]).await.unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].values, vec![3.0, 4.0]);

        // Count should still be 1
        assert_eq!(storage.count(&namespace).await.unwrap(), 1);
    }

    #[tokio::test]
    async fn test_index_storage() {
        let storage = ObjectStorage::memory().unwrap();
        let namespace = "test_index".to_string();

        // Initially no indexes
        assert!(!storage
            .index_exists(&namespace, IndexType::Hnsw)
            .await
            .unwrap());

        // Save an index
        let index_data = b"fake hnsw index data for testing".to_vec();
        storage
            .save_index(&namespace, IndexType::Hnsw, index_data.clone())
            .await
            .unwrap();

        // Check it exists
        assert!(storage
            .index_exists(&namespace, IndexType::Hnsw)
            .await
            .unwrap());
        assert!(!storage
            .index_exists(&namespace, IndexType::Pq)
            .await
            .unwrap());

        // Load it back
        let loaded = storage
            .load_index(&namespace, IndexType::Hnsw)
            .await
            .unwrap();
        assert!(loaded.is_some());
        assert_eq!(loaded.unwrap(), index_data);

        // Save another index type
        let pq_data = b"fake pq index data".to_vec();
        storage
            .save_index(&namespace, IndexType::Pq, pq_data)
            .await
            .unwrap();

        // List indexes
        let indexes = storage.list_indexes(&namespace).await.unwrap();
        assert_eq!(indexes.len(), 2);
        assert!(indexes.contains(&IndexType::Hnsw));
        assert!(indexes.contains(&IndexType::Pq));

        // Delete index
        let deleted = storage
            .delete_index(&namespace, IndexType::Hnsw)
            .await
            .unwrap();
        assert!(deleted);
        assert!(!storage
            .index_exists(&namespace, IndexType::Hnsw)
            .await
            .unwrap());

        // Delete non-existent
        let deleted = storage
            .delete_index(&namespace, IndexType::Hnsw)
            .await
            .unwrap();
        assert!(!deleted);

        // Load non-existent
        let loaded = storage
            .load_index(&namespace, IndexType::Hnsw)
            .await
            .unwrap();
        assert!(loaded.is_none());
    }
}