dakera-storage 0.10.1

//! Tiered Storage for Buffer
//!
//! Automatic data tiering based on access patterns:
//! - Hot tier: In-memory for frequently accessed data
//! - Warm tier: Local disk cache for recent data
//! - Cold tier: Object storage for infrequently accessed data

use async_trait::async_trait;
use common::{DakeraError, NamespaceId, Result, Vector, VectorId};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

use crate::traits::VectorStorage;

/// Tiered storage configuration
#[derive(Debug, Clone)]
pub struct TieredStorageConfig {
    /// Hot tier capacity (number of vectors)
    pub hot_tier_capacity: usize,
    /// Time before demoting from hot to warm
    pub hot_to_warm_threshold: Duration,
    /// Time before demoting from warm to cold
    pub warm_to_cold_threshold: Duration,
    /// Enable automatic tiering
    pub auto_tier_enabled: bool,
    /// Tier check interval
    pub tier_check_interval: Duration,
}

impl Default for TieredStorageConfig {
    fn default() -> Self {
        Self {
            hot_tier_capacity: 100_000,
            hot_to_warm_threshold: Duration::from_secs(3600), // 1 hour
            warm_to_cold_threshold: Duration::from_secs(86400), // 24 hours
            auto_tier_enabled: true,
            tier_check_interval: Duration::from_secs(300), // 5 minutes
        }
    }
}

/// Storage tier for a piece of data
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum StorageTier {
    /// In-memory, fastest access
    Hot,
    /// Local disk, fast access
    Warm,
    /// Object storage, slowest access
    Cold,
}

impl StorageTier {
    pub fn as_str(&self) -> &'static str {
        match self {
            StorageTier::Hot => "hot",
            StorageTier::Warm => "warm",
            StorageTier::Cold => "cold",
        }
    }
}

/// Access tracking for tiering decisions
#[derive(Debug, Clone)]
struct AccessInfo {
    /// Last access timestamp
    last_access: Instant,
    /// Total access count
    access_count: u64,
    /// Current tier
    tier: StorageTier,
}

impl Default for AccessInfo {
    fn default() -> Self {
        Self {
            last_access: Instant::now(),
            access_count: 0,
            tier: StorageTier::Hot,
        }
    }
}

/// Statistics for tiered storage
#[derive(Debug, Clone, Default)]
pub struct TieredStorageStats {
    /// Vectors in hot tier
    pub hot_count: u64,
    /// Vectors in warm tier
    pub warm_count: u64,
    /// Vectors in cold tier
    pub cold_count: u64,
    /// Hot tier hits
    pub hot_hits: u64,
    /// Warm tier hits
    pub warm_hits: u64,
    /// Cold tier hits
    pub cold_hits: u64,
    /// Promotions to hot tier
    pub promotions_to_hot: u64,
    /// Demotions to warm tier
    pub demotions_to_warm: u64,
    /// Demotions to cold tier
    pub demotions_to_cold: u64,
}

/// Tiered storage manager
pub struct TieredStorage<H, W, C> {
    /// Configuration
    config: TieredStorageConfig,
    /// Hot tier storage (in-memory)
    hot_storage: H,
    /// Warm tier storage (disk cache)
    warm_storage: W,
    /// Cold tier storage (object storage)
    cold_storage: C,
    /// Access tracking per vector
    access_info: RwLock<HashMap<(NamespaceId, VectorId), AccessInfo>>,
    /// Statistics
    stats: TieredStorageStatsInner,
}

struct TieredStorageStatsInner {
    hot_count: AtomicU64,
    warm_count: AtomicU64,
    cold_count: AtomicU64,
    hot_hits: AtomicU64,
    warm_hits: AtomicU64,
    cold_hits: AtomicU64,
    promotions_to_hot: AtomicU64,
    demotions_to_warm: AtomicU64,
    demotions_to_cold: AtomicU64,
}

impl Default for TieredStorageStatsInner {
    fn default() -> Self {
        Self {
            hot_count: AtomicU64::new(0),
            warm_count: AtomicU64::new(0),
            cold_count: AtomicU64::new(0),
            hot_hits: AtomicU64::new(0),
            warm_hits: AtomicU64::new(0),
            cold_hits: AtomicU64::new(0),
            promotions_to_hot: AtomicU64::new(0),
            demotions_to_warm: AtomicU64::new(0),
            demotions_to_cold: AtomicU64::new(0),
        }
    }
}

impl<H, W, C> TieredStorage<H, W, C>
where
    H: VectorStorage,
    W: VectorStorage,
    C: VectorStorage,
{
    /// Create a new tiered storage
    pub fn new(
        config: TieredStorageConfig,
        hot_storage: H,
        warm_storage: W,
        cold_storage: C,
    ) -> Self {
        Self {
            config,
            hot_storage,
            warm_storage,
            cold_storage,
            access_info: RwLock::new(HashMap::new()),
            stats: TieredStorageStatsInner::default(),
        }
    }

    /// Get the tiered storage configuration
    pub fn config(&self) -> &TieredStorageConfig {
        &self.config
    }

    /// Record access to a vector
    fn record_access(&self, namespace: &NamespaceId, id: &VectorId, tier: StorageTier) {
        let key = (namespace.clone(), id.clone());
        let mut access_map = self.access_info.write();

        let info = access_map.entry(key).or_default();
        info.last_access = Instant::now();
        info.access_count += 1;
        info.tier = tier;

        // Update hit counters
        match tier {
            StorageTier::Hot => self.stats.hot_hits.fetch_add(1, Ordering::Relaxed),
            StorageTier::Warm => self.stats.warm_hits.fetch_add(1, Ordering::Relaxed),
            StorageTier::Cold => self.stats.cold_hits.fetch_add(1, Ordering::Relaxed),
        };
    }

    /// Get the current tier for a vector
    fn get_tier(&self, namespace: &NamespaceId, id: &VectorId) -> Option<StorageTier> {
        let access_map = self.access_info.read();
        access_map
            .get(&(namespace.clone(), id.clone()))
            .map(|info| info.tier)
    }

    /// Promote a vector to a higher tier
    pub async fn promote(&self, namespace: &NamespaceId, id: &VectorId) -> Result<bool> {
        let current_tier = self.get_tier(namespace, id);

        match current_tier {
            Some(StorageTier::Warm) => {
                // Promote warm -> hot
                let vectors = self
                    .warm_storage
                    .get(namespace, std::slice::from_ref(id))
                    .await?;
                if !vectors.is_empty() {
                    self.hot_storage.upsert(namespace, vectors).await?;
                    self.warm_storage
                        .delete(namespace, std::slice::from_ref(id))
                        .await?;

                    self.update_tier(namespace, id, StorageTier::Hot);
                    self.stats.promotions_to_hot.fetch_add(1, Ordering::Relaxed);
                    self.stats.hot_count.fetch_add(1, Ordering::Relaxed);
                    self.stats.warm_count.fetch_sub(1, Ordering::Relaxed);

                    return Ok(true);
                }
            }
            Some(StorageTier::Cold) => {
                // Promote cold -> warm (or directly to hot if frequently accessed)
                let vectors = self
                    .cold_storage
                    .get(namespace, std::slice::from_ref(id))
                    .await?;
                if !vectors.is_empty() {
                    // Check if should go directly to hot based on access frequency
                    let should_be_hot = {
                        let access_map = self.access_info.read();
                        access_map
                            .get(&(namespace.clone(), id.clone()))
                            .map(|info| info.access_count > 10)
                            .unwrap_or(false)
                    };

                    if should_be_hot {
                        self.hot_storage.upsert(namespace, vectors).await?;
                        self.update_tier(namespace, id, StorageTier::Hot);
                        self.stats.promotions_to_hot.fetch_add(1, Ordering::Relaxed);
                        self.stats.hot_count.fetch_add(1, Ordering::Relaxed);
                    } else {
                        self.warm_storage.upsert(namespace, vectors).await?;
                        self.update_tier(namespace, id, StorageTier::Warm);
                        self.stats.warm_count.fetch_add(1, Ordering::Relaxed);
                    }
                    // Cold tier is the durable source of truth — never delete on promotion.
                    // The tier map tracks which tier is "active" for reads; cold remains
                    // as the persistent backup in case warm/hot are lost on restart.

                    return Ok(true);
                }
            }
            _ => {}
        }

        Ok(false)
    }

    /// Demote a vector to a lower tier
    pub async fn demote(&self, namespace: &NamespaceId, id: &VectorId) -> Result<bool> {
        let current_tier = self.get_tier(namespace, id);

        match current_tier {
            Some(StorageTier::Hot) => {
                // Demote hot -> warm
                let vectors = self
                    .hot_storage
                    .get(namespace, std::slice::from_ref(id))
                    .await?;
                if !vectors.is_empty() {
                    self.warm_storage.upsert(namespace, vectors).await?;
                    self.hot_storage
                        .delete(namespace, std::slice::from_ref(id))
                        .await?;

                    self.update_tier(namespace, id, StorageTier::Warm);
                    self.stats.demotions_to_warm.fetch_add(1, Ordering::Relaxed);
                    self.stats.hot_count.fetch_sub(1, Ordering::Relaxed);
                    self.stats.warm_count.fetch_add(1, Ordering::Relaxed);

                    return Ok(true);
                }
            }
            Some(StorageTier::Warm) => {
                // Demote warm -> cold
                let vectors = self
                    .warm_storage
                    .get(namespace, std::slice::from_ref(id))
                    .await?;
                if !vectors.is_empty() {
                    self.cold_storage.upsert(namespace, vectors).await?;
                    self.warm_storage
                        .delete(namespace, std::slice::from_ref(id))
                        .await?;

                    self.update_tier(namespace, id, StorageTier::Cold);
                    self.stats.demotions_to_cold.fetch_add(1, Ordering::Relaxed);
                    self.stats.warm_count.fetch_sub(1, Ordering::Relaxed);
                    self.stats.cold_count.fetch_add(1, Ordering::Relaxed);

                    return Ok(true);
                }
            }
            _ => {}
        }

        Ok(false)
    }

    /// Update tier tracking
    fn update_tier(&self, namespace: &NamespaceId, id: &VectorId, tier: StorageTier) {
        let mut access_map = self.access_info.write();
        let key = (namespace.clone(), id.clone());
        let info = access_map.entry(key).or_default();
        info.tier = tier;
    }

    /// Run automatic tiering based on access patterns
    pub async fn run_auto_tiering(&self) -> Result<TieringResult> {
        if !self.config.auto_tier_enabled {
            return Ok(TieringResult::default());
        }

        let now = Instant::now();
        let mut to_demote_to_warm = Vec::new();
        let mut to_demote_to_cold = Vec::new();

        // Collect vectors to demote
        {
            let access_map = self.access_info.read();
            for ((namespace, id), info) in access_map.iter() {
                let elapsed = now.duration_since(info.last_access);

                match info.tier {
                    StorageTier::Hot if elapsed > self.config.hot_to_warm_threshold => {
                        to_demote_to_warm.push((namespace.clone(), id.clone()));
                    }
                    StorageTier::Warm if elapsed > self.config.warm_to_cold_threshold => {
                        to_demote_to_cold.push((namespace.clone(), id.clone()));
                    }
                    _ => {}
                }
            }
        }

        // Execute demotions
        let mut demoted_to_warm = 0;
        let mut demoted_to_cold = 0;

        for (namespace, id) in to_demote_to_warm {
            if self.demote(&namespace, &id).await? {
                demoted_to_warm += 1;
            }
        }

        for (namespace, id) in to_demote_to_cold {
            if self.demote(&namespace, &id).await? {
                demoted_to_cold += 1;
            }
        }

        Ok(TieringResult {
            demoted_to_warm,
            demoted_to_cold,
            promoted_to_hot: 0,
            promoted_to_warm: 0,
        })
    }

    /// Get storage statistics
    pub fn stats(&self) -> TieredStorageStats {
        TieredStorageStats {
            hot_count: self.stats.hot_count.load(Ordering::Relaxed),
            warm_count: self.stats.warm_count.load(Ordering::Relaxed),
            cold_count: self.stats.cold_count.load(Ordering::Relaxed),
            hot_hits: self.stats.hot_hits.load(Ordering::Relaxed),
            warm_hits: self.stats.warm_hits.load(Ordering::Relaxed),
            cold_hits: self.stats.cold_hits.load(Ordering::Relaxed),
            promotions_to_hot: self.stats.promotions_to_hot.load(Ordering::Relaxed),
            demotions_to_warm: self.stats.demotions_to_warm.load(Ordering::Relaxed),
            demotions_to_cold: self.stats.demotions_to_cold.load(Ordering::Relaxed),
        }
    }

    /// Get tier distribution by namespace
    pub fn tier_distribution(&self, namespace: &NamespaceId) -> TierDistribution {
        let access_map = self.access_info.read();
        let mut hot = 0u64;
        let mut warm = 0u64;
        let mut cold = 0u64;

        for ((ns, _), info) in access_map.iter() {
            if ns == namespace {
                match info.tier {
                    StorageTier::Hot => hot += 1,
                    StorageTier::Warm => warm += 1,
                    StorageTier::Cold => cold += 1,
                }
            }
        }

        TierDistribution { hot, warm, cold }
    }
}

/// Result of automatic tiering
#[derive(Debug, Clone, Default)]
pub struct TieringResult {
    /// Vectors demoted to warm tier
    pub demoted_to_warm: u64,
    /// Vectors demoted to cold tier
    pub demoted_to_cold: u64,
    /// Vectors promoted to hot tier
    pub promoted_to_hot: u64,
    /// Vectors promoted to warm tier
    pub promoted_to_warm: u64,
}

/// Tier distribution for a namespace
#[derive(Debug, Clone)]
pub struct TierDistribution {
    pub hot: u64,
    pub warm: u64,
    pub cold: u64,
}

#[async_trait]
impl<H, W, C> VectorStorage for TieredStorage<H, W, C>
where
    H: VectorStorage,
    W: VectorStorage,
    C: VectorStorage + Clone + Send + Sync + 'static,
{
    async fn upsert(&self, namespace: &NamespaceId, vectors: Vec<Vector>) -> Result<usize> {
        // Track access info before moving vectors into storage
        let ids: Vec<VectorId> = vectors.iter().map(|v| v.id.clone()).collect();

        // Clone vectors for the background cold-tier flush before consuming them.
        let cold_vectors = vectors.clone();

        // Write hot tier first — lowest latency path, immediately available for reads.
        let count = self.hot_storage.upsert(namespace, vectors).await?;

        // Flush to cold tier (S3) in the background — failure is logged but not fatal.
        // The hot tier (RocksDB/in-memory) is the primary durable read path.
        let cold = self.cold_storage.clone();
        let cold_ns = namespace.clone();
        tokio::spawn(async move {
            if let Err(e) = cold.ensure_namespace(&cold_ns).await {
                tracing::error!(
                    error = %e,
                    namespace = %cold_ns,
                    "Cold tier namespace ensure failed (S3 flush aborted)"
                );
                return;
            }
            if let Err(e) = cold.upsert(&cold_ns, cold_vectors).await {
                tracing::error!(
                    error = %e,
                    namespace = %cold_ns,
                    "Cold tier S3 flush failed — data is durable in hot tier"
                );
            }
        });

        for id in &ids {
            self.update_tier(namespace, id, StorageTier::Hot);
            self.record_access(namespace, id, StorageTier::Hot);
        }

        self.stats
            .hot_count
            .fetch_add(count as u64, Ordering::Relaxed);
        Ok(count)
    }

    async fn get(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<Vec<Vector>> {
        let mut results = Vec::with_capacity(ids.len());
        let mut remaining_ids: Vec<VectorId> = ids.to_vec();

        // Try hot tier first (NamespaceNotFound is normal after restart when hot cache is empty)
        let hot_results = match self.hot_storage.get(namespace, &remaining_ids).await {
            Ok(v) => v,
            Err(DakeraError::NamespaceNotFound(_)) => vec![],
            Err(e) => return Err(e),
        };
        for v in &hot_results {
            self.record_access(namespace, &v.id, StorageTier::Hot);
        }

        // Remove found IDs
        let found_ids: std::collections::HashSet<_> = hot_results.iter().map(|v| &v.id).collect();
        remaining_ids.retain(|id| !found_ids.contains(id));
        results.extend(hot_results);

        if remaining_ids.is_empty() {
            return Ok(results);
        }

        // Try warm tier (NamespaceNotFound is normal after restart when warm cache is empty)
        let warm_results = match self.warm_storage.get(namespace, &remaining_ids).await {
            Ok(v) => v,
            Err(common::DakeraError::NamespaceNotFound(_)) => vec![],
            Err(e) => return Err(e),
        };
        for v in &warm_results {
            self.record_access(namespace, &v.id, StorageTier::Warm);
        }

        let found_ids: std::collections::HashSet<_> = warm_results.iter().map(|v| &v.id).collect();
        remaining_ids.retain(|id| !found_ids.contains(id));
        results.extend(warm_results);

        if remaining_ids.is_empty() {
            return Ok(results);
        }

        // Try cold tier
        let cold_results = match self.cold_storage.get(namespace, &remaining_ids).await {
            Ok(v) => v,
            Err(DakeraError::NamespaceNotFound(_)) => vec![],
            Err(e) => return Err(e),
        };
        for v in &cold_results {
            self.record_access(namespace, &v.id, StorageTier::Cold);
        }
        results.extend(cold_results);

        Ok(results)
    }

    async fn get_all(&self, namespace: &NamespaceId) -> Result<Vec<Vector>> {
        let mut seen = std::collections::HashSet::new();
        let mut results = Vec::new();

        // Helper: treat NamespaceNotFound as empty — normal for hot/warm after restart.
        let tier_get_all = |res: common::Result<Vec<Vector>>| -> common::Result<Vec<Vector>> {
            match res {
                Ok(v) => Ok(v),
                Err(common::DakeraError::NamespaceNotFound(_)) => Ok(vec![]),
                Err(e) => Err(e),
            }
        };

        // Gather from all tiers, preferring hot over warm over cold.
        // Deduplicate by vector ID since write-through means a vector
        // can exist in both hot and cold simultaneously.
        for v in tier_get_all(self.hot_storage.get_all(namespace).await)? {
            if seen.insert(v.id.clone()) {
                results.push(v);
            }
        }
        for v in tier_get_all(self.warm_storage.get_all(namespace).await)? {
            if seen.insert(v.id.clone()) {
                results.push(v);
            }
        }
        for v in tier_get_all(self.cold_storage.get_all(namespace).await)? {
            if seen.insert(v.id.clone()) {
                results.push(v);
            }
        }

        Ok(results)
    }

    async fn delete(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<usize> {
        let mut deleted = 0;

        // Delete from all tiers. Tolerate NamespaceNotFound from individual tiers
        // since data may only reside in a subset of tiers (e.g. cold but not hot).
        match self.hot_storage.delete(namespace, ids).await {
            Ok(n) => deleted += n,
            Err(DakeraError::NamespaceNotFound(_)) => {}
            Err(e) => return Err(e),
        }
        match self.warm_storage.delete(namespace, ids).await {
            Ok(n) => deleted += n,
            Err(DakeraError::NamespaceNotFound(_)) => {}
            Err(e) => return Err(e),
        }
        match self.cold_storage.delete(namespace, ids).await {
            Ok(n) => deleted += n,
            Err(DakeraError::NamespaceNotFound(_)) => {}
            Err(e) => return Err(e),
        }

        // Remove from tracking
        {
            let mut access_map = self.access_info.write();
            for id in ids {
                access_map.remove(&(namespace.clone(), id.clone()));
            }
        }

        Ok(deleted)
    }

    async fn namespace_exists(&self, namespace: &NamespaceId) -> Result<bool> {
        // Check any tier
        Ok(self.hot_storage.namespace_exists(namespace).await?
            || self.warm_storage.namespace_exists(namespace).await?
            || self.cold_storage.namespace_exists(namespace).await?)
    }

    async fn ensure_namespace(&self, namespace: &NamespaceId) -> Result<()> {
        // Ensure in all tiers
        self.hot_storage.ensure_namespace(namespace).await?;
        self.warm_storage.ensure_namespace(namespace).await?;
        self.cold_storage.ensure_namespace(namespace).await?;
        Ok(())
    }

    async fn count(&self, namespace: &NamespaceId) -> Result<usize> {
        // With write-through, cold tier is the source of truth for total count.
        // Hot/warm are caches that hold subsets of the same data.
        // Use cold count as the baseline, then add any vectors that are ONLY
        // in hot or warm (shouldn't happen with write-through, but safe).
        let cold = self.cold_storage.count(namespace).await?;
        if cold > 0 {
            return Ok(cold);
        }
        // Fallback: if cold is empty, count from hot + warm (non-tiered data)
        let hot = self.hot_storage.count(namespace).await?;
        let warm = self.warm_storage.count(namespace).await?;
        Ok(hot + warm)
    }

    async fn dimension(&self, namespace: &NamespaceId) -> Result<Option<usize>> {
        // Check hot first, then warm, then cold
        if let Some(dim) = self.hot_storage.dimension(namespace).await? {
            return Ok(Some(dim));
        }
        if let Some(dim) = self.warm_storage.dimension(namespace).await? {
            return Ok(Some(dim));
        }
        self.cold_storage.dimension(namespace).await
    }

    async fn list_namespaces(&self) -> Result<Vec<NamespaceId>> {
        let mut namespaces = std::collections::HashSet::new();

        namespaces.extend(self.hot_storage.list_namespaces().await?);
        namespaces.extend(self.warm_storage.list_namespaces().await?);
        namespaces.extend(self.cold_storage.list_namespaces().await?);

        Ok(namespaces.into_iter().collect())
    }

    async fn delete_namespace(&self, namespace: &NamespaceId) -> Result<bool> {
        // Delete from all tiers
        let hot_deleted = self.hot_storage.delete_namespace(namespace).await?;
        let warm_deleted = self.warm_storage.delete_namespace(namespace).await?;
        let cold_deleted = self.cold_storage.delete_namespace(namespace).await?;

        // Remove from access tracking
        {
            let mut access_map = self.access_info.write();
            access_map.retain(|(ns, _), _| ns != namespace);
        }

        Ok(hot_deleted || warm_deleted || cold_deleted)
    }

    async fn cleanup_expired(&self, namespace: &NamespaceId) -> Result<usize> {
        // Cleanup from all tiers
        let mut total = 0;
        total += self.hot_storage.cleanup_expired(namespace).await?;
        total += self.warm_storage.cleanup_expired(namespace).await?;
        total += self.cold_storage.cleanup_expired(namespace).await?;
        Ok(total)
    }

    async fn cleanup_all_expired(&self) -> Result<usize> {
        // Cleanup from all tiers
        let mut total = 0;
        total += self.hot_storage.cleanup_all_expired().await?;
        total += self.warm_storage.cleanup_all_expired().await?;
        total += self.cold_storage.cleanup_all_expired().await?;
        Ok(total)
    }
}

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

    fn create_test_vector(id: &str, dim: usize) -> Vector {
        Vector {
            id: id.to_string(),
            values: vec![1.0; dim],
            metadata: None,
            ttl_seconds: None,
            expires_at: None,
        }
    }

    #[tokio::test]
    async fn test_tiered_storage_basic() {
        let config = TieredStorageConfig::default();
        let storage = TieredStorage::new(
            config,
            InMemoryStorage::new(),
            InMemoryStorage::new(),
            InMemoryStorage::new(),
        );

        let namespace = "test".to_string();
        storage.ensure_namespace(&namespace).await.unwrap();

        // Upsert goes to hot tier
        let vectors = vec![create_test_vector("v1", 4)];
        let count = storage.upsert(&namespace, vectors).await.unwrap();
        assert_eq!(count, 1);

        // Get should find in hot tier
        let results = storage.get(&namespace, &["v1".to_string()]).await.unwrap();
        assert_eq!(results.len(), 1);

        let stats = storage.stats();
        assert_eq!(stats.hot_hits, 2); // One from upsert record_access, one from get
    }

    #[tokio::test]
    async fn test_tiered_storage_promotion_demotion() {
        let config = TieredStorageConfig::default();
        let storage = TieredStorage::new(
            config,
            InMemoryStorage::new(),
            InMemoryStorage::new(),
            InMemoryStorage::new(),
        );

        let namespace = "test".to_string();
        storage.ensure_namespace(&namespace).await.unwrap();

        // Add to hot tier
        storage
            .upsert(&namespace, vec![create_test_vector("v1", 4)])
            .await
            .unwrap();

        // Verify in hot
        assert_eq!(
            storage.get_tier(&namespace, &"v1".to_string()),
            Some(StorageTier::Hot)
        );

        // Demote to warm
        let demoted = storage.demote(&namespace, &"v1".to_string()).await.unwrap();
        assert!(demoted);
        assert_eq!(
            storage.get_tier(&namespace, &"v1".to_string()),
            Some(StorageTier::Warm)
        );

        // Demote to cold
        let demoted = storage.demote(&namespace, &"v1".to_string()).await.unwrap();
        assert!(demoted);
        assert_eq!(
            storage.get_tier(&namespace, &"v1".to_string()),
            Some(StorageTier::Cold)
        );

        // Still accessible
        let results = storage.get(&namespace, &["v1".to_string()]).await.unwrap();
        assert_eq!(results.len(), 1);

        let stats = storage.stats();
        assert_eq!(stats.demotions_to_warm, 1);
        assert_eq!(stats.demotions_to_cold, 1);
    }

    #[tokio::test]
    async fn test_tiered_storage_multi_tier_get() {
        let config = TieredStorageConfig::default();
        let storage = TieredStorage::new(
            config,
            InMemoryStorage::new(),
            InMemoryStorage::new(),
            InMemoryStorage::new(),
        );

        let namespace = "test".to_string();
        storage.ensure_namespace(&namespace).await.unwrap();

        // Add vectors and demote some
        for i in 0..3 {
            storage
                .upsert(&namespace, vec![create_test_vector(&format!("v{}", i), 4)])
                .await
                .unwrap();
        }

        // v0 stays hot, v1 goes warm, v2 goes cold
        storage.demote(&namespace, &"v1".to_string()).await.unwrap();
        storage.demote(&namespace, &"v2".to_string()).await.unwrap();
        storage.demote(&namespace, &"v2".to_string()).await.unwrap();

        // Get all at once
        let ids: Vec<_> = (0..3).map(|i| format!("v{}", i)).collect();
        let results = storage.get(&namespace, &ids).await.unwrap();
        assert_eq!(results.len(), 3);
    }

    #[tokio::test]
    async fn test_tier_distribution() {
        let config = TieredStorageConfig::default();
        let storage = TieredStorage::new(
            config,
            InMemoryStorage::new(),
            InMemoryStorage::new(),
            InMemoryStorage::new(),
        );

        let namespace = "test".to_string();
        storage.ensure_namespace(&namespace).await.unwrap();

        // Add 5 vectors
        for i in 0..5 {
            storage
                .upsert(&namespace, vec![create_test_vector(&format!("v{}", i), 4)])
                .await
                .unwrap();
        }

        // Demote 2 to warm, 1 to cold
        storage.demote(&namespace, &"v3".to_string()).await.unwrap();
        storage.demote(&namespace, &"v4".to_string()).await.unwrap();
        storage.demote(&namespace, &"v4".to_string()).await.unwrap();

        let dist = storage.tier_distribution(&namespace);
        assert_eq!(dist.hot, 3);
        assert_eq!(dist.warm, 1);
        assert_eq!(dist.cold, 1);
    }
}