storage/

compaction.rs

//! Compaction Module for Object Storage
//!
//! This module provides segment management and compaction for object storage backends.
//! It helps maintain storage efficiency by:
//!
//! - Organizing vectors into segments for batch I/O
//! - Merging small segments into larger ones
//! - Garbage collecting deleted vectors
//! - Running background compaction based on configurable thresholds
//!
//! # Architecture
//!
//! ```text
//! ┌─────────────────────────────────────────────────────┐
//! │                  CompactionManager                   │
//! │  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐ │
//! │  │  Segment 1  │  │  Segment 2  │  │  Segment N  │ │
//! │  │  (1000 vec) │  │  (500 vec)  │  │  (750 vec)  │ │
//! │  └─────────────┘  └─────────────┘  └─────────────┘ │
//! │         │                │                │         │
//! │         └────────────────┼────────────────┘         │
//! │                          ▼                          │
//! │                   Compaction Job                    │
//! │                          │                          │
//! │                          ▼                          │
//! │              ┌─────────────────────┐               │
//! │              │   Merged Segment    │               │
//! │              │    (2250 vec)       │               │
//! │              └─────────────────────┘               │
//! └─────────────────────────────────────────────────────┘
//! ```

use common::{NamespaceId, Vector, VectorId};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

// ============================================================================
// Compaction Configuration
// ============================================================================

/// Configuration for compaction operations
#[derive(Debug, Clone)]
pub struct CompactionConfig {
    /// Target segment size (number of vectors)
    pub target_segment_size: usize,
    /// Minimum segment size before considering for merge
    pub min_segment_size: usize,
    /// Maximum segments to merge in one operation
    pub max_merge_segments: usize,
    /// Garbage ratio threshold to trigger compaction (0.0 - 1.0)
    pub garbage_threshold: f32,
    /// Enable automatic background compaction
    pub auto_compact: bool,
    /// Interval between compaction checks (seconds)
    pub compaction_interval_secs: u64,
    /// Maximum concurrent compaction jobs
    pub max_concurrent_jobs: usize,
    /// Preserve tombstones for this duration (seconds) for consistency
    pub tombstone_ttl_secs: u64,
}

impl Default for CompactionConfig {
    fn default() -> Self {
        Self {
            target_segment_size: 10_000,
            min_segment_size: 1_000,
            max_merge_segments: 5,
            garbage_threshold: 0.3, // Compact when 30% is garbage
            auto_compact: true,
            compaction_interval_secs: 300, // 5 minutes
            max_concurrent_jobs: 2,
            tombstone_ttl_secs: 3600, // 1 hour
        }
    }
}

impl CompactionConfig {
    /// Create a new config
    pub fn new() -> Self {
        Self::default()
    }

    /// Set target segment size
    pub fn with_target_size(mut self, size: usize) -> Self {
        self.target_segment_size = size;
        self
    }

    /// Set minimum segment size
    pub fn with_min_size(mut self, size: usize) -> Self {
        self.min_segment_size = size;
        self
    }

    /// Set garbage threshold
    pub fn with_garbage_threshold(mut self, threshold: f32) -> Self {
        self.garbage_threshold = threshold.clamp(0.0, 1.0);
        self
    }

    /// Disable automatic compaction
    pub fn without_auto_compact(mut self) -> Self {
        self.auto_compact = false;
        self
    }

    /// Set compaction interval
    pub fn with_interval(mut self, secs: u64) -> Self {
        self.compaction_interval_secs = secs;
        self
    }
}

// ============================================================================
// Segment Types
// ============================================================================

/// Unique identifier for a segment
pub type SegmentId = String;

/// Segment state
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum SegmentState {
    /// Segment is active and accepting writes
    Active,
    /// Segment is sealed and read-only
    Sealed,
    /// Segment is being compacted
    Compacting,
    /// Segment is marked for deletion
    Tombstone,
}

/// Metadata for a segment
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SegmentMetadata {
    /// Segment ID
    pub id: SegmentId,
    /// Namespace this segment belongs to
    pub namespace: NamespaceId,
    /// Current state
    pub state: SegmentState,
    /// Number of live vectors
    pub live_count: usize,
    /// Number of deleted vectors (garbage)
    pub deleted_count: usize,
    /// Total size in bytes (estimated)
    pub size_bytes: usize,
    /// Creation timestamp
    pub created_at: u64,
    /// Last modified timestamp
    pub updated_at: u64,
    /// Minimum vector ID in segment (for range queries)
    pub min_id: Option<VectorId>,
    /// Maximum vector ID in segment (for range queries)
    pub max_id: Option<VectorId>,
    /// Level in LSM-like hierarchy (0 = most recent)
    pub level: u32,
}

impl SegmentMetadata {
    /// Create new segment metadata
    pub fn new(id: SegmentId, namespace: NamespaceId) -> Self {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        Self {
            id,
            namespace,
            state: SegmentState::Active,
            live_count: 0,
            deleted_count: 0,
            size_bytes: 0,
            created_at: now,
            updated_at: now,
            min_id: None,
            max_id: None,
            level: 0,
        }
    }

    /// Calculate garbage ratio
    pub fn garbage_ratio(&self) -> f32 {
        let total = self.live_count + self.deleted_count;
        if total == 0 {
            0.0
        } else {
            self.deleted_count as f32 / total as f32
        }
    }

    /// Check if segment is empty
    pub fn is_empty(&self) -> bool {
        self.live_count == 0
    }

    /// Total vector count (live + deleted)
    pub fn total_count(&self) -> usize {
        self.live_count + self.deleted_count
    }

    /// Check if segment should be compacted based on garbage ratio
    pub fn needs_compaction(&self, threshold: f32) -> bool {
        self.garbage_ratio() >= threshold
    }

    /// Update the ID range
    pub fn update_id_range(&mut self, id: &VectorId) {
        match &self.min_id {
            None => self.min_id = Some(id.clone()),
            Some(min) if id < min => self.min_id = Some(id.clone()),
            _ => {}
        }
        match &self.max_id {
            None => self.max_id = Some(id.clone()),
            Some(max) if id > max => self.max_id = Some(id.clone()),
            _ => {}
        }
    }
}

/// A segment containing vectors
#[derive(Debug, Clone)]
pub struct Segment {
    /// Segment metadata
    pub metadata: SegmentMetadata,
    /// Vectors in this segment
    vectors: HashMap<VectorId, Vector>,
    /// Set of deleted vector IDs (tombstones)
    tombstones: HashSet<VectorId>,
}

impl Segment {
    /// Create a new segment
    pub fn new(id: SegmentId, namespace: NamespaceId) -> Self {
        Self {
            metadata: SegmentMetadata::new(id, namespace),
            vectors: HashMap::new(),
            tombstones: HashSet::new(),
        }
    }

    /// Add a vector to the segment
    pub fn add(&mut self, vector: Vector) {
        let size = estimate_vector_size(&vector);
        self.metadata.update_id_range(&vector.id);

        // Remove from tombstones if re-adding
        if self.tombstones.remove(&vector.id) {
            self.metadata.deleted_count = self.metadata.deleted_count.saturating_sub(1);
        }

        // Check if updating existing
        if self.vectors.contains_key(&vector.id) {
            // Update existing - size change
            let old_size = self
                .vectors
                .get(&vector.id)
                .map(estimate_vector_size)
                .unwrap_or(0);
            self.metadata.size_bytes = self.metadata.size_bytes.saturating_sub(old_size);
        } else {
            self.metadata.live_count += 1;
        }

        self.metadata.size_bytes += size;
        self.metadata.updated_at = current_timestamp();
        self.vectors.insert(vector.id.clone(), vector);
    }

    /// Delete a vector from the segment (mark as tombstone)
    pub fn delete(&mut self, id: &VectorId) -> bool {
        if let Some(vector) = self.vectors.remove(id) {
            let size = estimate_vector_size(&vector);
            self.metadata.live_count = self.metadata.live_count.saturating_sub(1);
            self.metadata.deleted_count += 1;
            self.metadata.size_bytes = self.metadata.size_bytes.saturating_sub(size);
            self.metadata.updated_at = current_timestamp();
            self.tombstones.insert(id.clone());
            true
        } else {
            false
        }
    }

    /// Get a vector by ID
    pub fn get(&self, id: &VectorId) -> Option<&Vector> {
        if self.tombstones.contains(id) {
            None
        } else {
            self.vectors.get(id)
        }
    }

    /// Check if a vector exists (and is not deleted)
    pub fn contains(&self, id: &VectorId) -> bool {
        !self.tombstones.contains(id) && self.vectors.contains_key(id)
    }

    /// Check if a vector is tombstoned
    pub fn is_tombstoned(&self, id: &VectorId) -> bool {
        self.tombstones.contains(id)
    }

    /// Get all live vectors
    pub fn live_vectors(&self) -> impl Iterator<Item = &Vector> {
        self.vectors
            .values()
            .filter(|v| !self.tombstones.contains(&v.id))
    }

    /// Get all vector IDs (including tombstoned)
    pub fn all_ids(&self) -> impl Iterator<Item = &VectorId> {
        self.vectors.keys()
    }

    /// Get tombstone IDs
    pub fn tombstone_ids(&self) -> impl Iterator<Item = &VectorId> {
        self.tombstones.iter()
    }

    /// Seal the segment (make read-only)
    pub fn seal(&mut self) {
        self.metadata.state = SegmentState::Sealed;
        self.metadata.updated_at = current_timestamp();
    }

    /// Mark as compacting
    pub fn mark_compacting(&mut self) {
        self.metadata.state = SegmentState::Compacting;
        self.metadata.updated_at = current_timestamp();
    }

    /// Mark for deletion
    pub fn mark_tombstone(&mut self) {
        self.metadata.state = SegmentState::Tombstone;
        self.metadata.updated_at = current_timestamp();
    }
}

// ============================================================================
// Compaction Operations
// ============================================================================

/// Result of a compaction operation
#[derive(Debug, Clone)]
pub struct CompactionResult {
    /// Segments that were merged
    pub merged_segments: Vec<SegmentId>,
    /// New segment created (if any)
    pub new_segment: Option<SegmentId>,
    /// Number of vectors compacted
    pub vectors_compacted: usize,
    /// Number of tombstones removed
    pub tombstones_removed: usize,
    /// Bytes reclaimed
    pub bytes_reclaimed: usize,
    /// Duration of compaction (milliseconds)
    pub duration_ms: u64,
}

/// Compaction job status
#[derive(Debug, Clone)]
pub struct CompactionJob {
    /// Job ID
    pub id: String,
    /// Namespace being compacted
    pub namespace: NamespaceId,
    /// Segments being merged
    pub source_segments: Vec<SegmentId>,
    /// Target segment ID
    pub target_segment: SegmentId,
    /// Job state
    pub state: CompactionJobState,
    /// Progress (0.0 - 1.0)
    pub progress: f32,
    /// Start time
    pub started_at: u64,
    /// Completion time (if finished)
    pub completed_at: Option<u64>,
    /// Error message (if failed)
    pub error: Option<String>,
}

/// State of a compaction job
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompactionJobState {
    /// Job is queued
    Pending,
    /// Job is running
    Running,
    /// Job completed successfully
    Completed,
    /// Job failed
    Failed,
    /// Job was cancelled
    Cancelled,
}

// ============================================================================
// Namespace Segment Manager
// ============================================================================

/// Manages segments for a namespace
pub struct NamespaceSegmentManager {
    /// Namespace ID
    namespace: NamespaceId,
    /// Configuration
    config: CompactionConfig,
    /// Active segment (for writes)
    active_segment: RwLock<Option<Segment>>,
    /// Sealed segments
    sealed_segments: RwLock<HashMap<SegmentId, Segment>>,
    /// Vector ID to segment ID mapping
    vector_index: RwLock<HashMap<VectorId, SegmentId>>,
    /// Segment counter for ID generation
    segment_counter: AtomicU64,
    /// Statistics
    stats: CompactionStats,
}

impl NamespaceSegmentManager {
    /// Create a new namespace segment manager
    pub fn new(namespace: NamespaceId, config: CompactionConfig) -> Self {
        Self {
            namespace,
            config,
            active_segment: RwLock::new(None),
            sealed_segments: RwLock::new(HashMap::new()),
            vector_index: RwLock::new(HashMap::new()),
            segment_counter: AtomicU64::new(0),
            stats: CompactionStats::default(),
        }
    }

    /// Generate a new segment ID
    fn generate_segment_id(&self) -> SegmentId {
        let counter = self.segment_counter.fetch_add(1, Ordering::SeqCst);
        format!("{}_{:016x}", self.namespace, counter)
    }

    /// Get or create the active segment
    fn ensure_active_segment(&self) -> SegmentId {
        let mut active = self.active_segment.write();
        if active.is_none() {
            let id = self.generate_segment_id();
            *active = Some(Segment::new(id.clone(), self.namespace.clone()));
            return id;
        }

        let segment = active.as_ref().unwrap();

        // Check if we need to seal and create new segment
        if segment.metadata.live_count >= self.config.target_segment_size {
            // Seal current segment
            let mut sealed = active.take().unwrap();
            sealed.seal();
            let sealed_id = sealed.metadata.id.clone();

            // Move to sealed segments
            self.sealed_segments.write().insert(sealed_id, sealed);

            // Create new active segment
            let id = self.generate_segment_id();
            *active = Some(Segment::new(id.clone(), self.namespace.clone()));
            return id;
        }

        segment.metadata.id.clone()
    }

    /// Add a vector
    pub fn add(&self, vector: Vector) {
        let segment_id = self.ensure_active_segment();

        // Update vector index
        self.vector_index
            .write()
            .insert(vector.id.clone(), segment_id.clone());

        // Add to active segment
        if let Some(ref mut segment) = *self.active_segment.write() {
            segment.add(vector);
        }

        self.stats.vectors_written.fetch_add(1, Ordering::Relaxed);
    }

    /// Get a vector by ID
    pub fn get(&self, id: &VectorId) -> Option<Vector> {
        // First check vector index
        let segment_id = self.vector_index.read().get(id)?.clone();

        // Check active segment
        if let Some(ref segment) = *self.active_segment.read() {
            if segment.metadata.id == segment_id {
                return segment.get(id).cloned();
            }
        }

        // Check sealed segments
        self.sealed_segments
            .read()
            .get(&segment_id)
            .and_then(|s| s.get(id).cloned())
    }

    /// Delete a vector
    pub fn delete(&self, id: &VectorId) -> bool {
        // Find the segment
        let segment_id = match self.vector_index.read().get(id) {
            Some(id) => id.clone(),
            None => return false,
        };

        // Try active segment
        if let Some(ref mut segment) = *self.active_segment.write() {
            if segment.metadata.id == segment_id && segment.delete(id) {
                self.stats.vectors_deleted.fetch_add(1, Ordering::Relaxed);
                return true;
            }
        }

        // Try sealed segments
        if let Some(segment) = self.sealed_segments.write().get_mut(&segment_id) {
            if segment.delete(id) {
                self.stats.vectors_deleted.fetch_add(1, Ordering::Relaxed);
                return true;
            }
        }

        false
    }

    /// Get all vectors
    pub fn get_all(&self) -> Vec<Vector> {
        let mut result = Vec::new();

        // Active segment
        if let Some(ref segment) = *self.active_segment.read() {
            result.extend(segment.live_vectors().cloned());
        }

        // Sealed segments
        for segment in self.sealed_segments.read().values() {
            result.extend(segment.live_vectors().cloned());
        }

        result
    }

    /// Get segments that need compaction
    pub fn segments_needing_compaction(&self) -> Vec<SegmentMetadata> {
        let sealed = self.sealed_segments.read();
        sealed
            .values()
            .filter(|s| {
                s.metadata.state == SegmentState::Sealed
                    && s.metadata.needs_compaction(self.config.garbage_threshold)
            })
            .map(|s| s.metadata.clone())
            .collect()
    }

    /// Get small segments that could be merged
    pub fn small_segments(&self) -> Vec<SegmentMetadata> {
        let sealed = self.sealed_segments.read();
        sealed
            .values()
            .filter(|s| {
                s.metadata.state == SegmentState::Sealed
                    && s.metadata.live_count < self.config.min_segment_size
            })
            .map(|s| s.metadata.clone())
            .collect()
    }

    /// Perform compaction on specified segments
    pub fn compact(&self, segment_ids: &[SegmentId]) -> Option<CompactionResult> {
        if segment_ids.is_empty() {
            return None;
        }

        let start = std::time::Instant::now();
        let mut sealed = self.sealed_segments.write();
        let mut vector_index = self.vector_index.write();

        // Collect segments to merge
        let mut segments_to_merge: Vec<Segment> = Vec::new();
        for id in segment_ids {
            if let Some(mut segment) = sealed.remove(id) {
                segment.mark_compacting();
                segments_to_merge.push(segment);
            }
        }

        if segments_to_merge.is_empty() {
            return None;
        }

        // Create new merged segment
        let new_id = self.generate_segment_id();
        let mut new_segment = Segment::new(new_id.clone(), self.namespace.clone());
        new_segment.metadata.level = segments_to_merge
            .iter()
            .map(|s| s.metadata.level)
            .max()
            .unwrap_or(0)
            + 1;

        let mut vectors_compacted = 0;
        let mut tombstones_removed = 0;
        let mut bytes_reclaimed = 0;

        // Merge all live vectors into new segment
        for segment in &segments_to_merge {
            bytes_reclaimed += segment.metadata.size_bytes;
            tombstones_removed += segment.tombstones.len();

            for vector in segment.live_vectors() {
                new_segment.add(vector.clone());
                vector_index.insert(vector.id.clone(), new_id.clone());
                vectors_compacted += 1;
            }

            // Remove tombstones from vector index
            for tombstone_id in segment.tombstone_ids() {
                vector_index.remove(tombstone_id);
            }
        }

        // Seal new segment if it's large enough
        if new_segment.metadata.live_count >= self.config.min_segment_size {
            new_segment.seal();
        }

        // Calculate actual bytes (subtract new segment size from reclaimed)
        bytes_reclaimed = bytes_reclaimed.saturating_sub(new_segment.metadata.size_bytes);

        // Store new segment
        let merged_ids: Vec<_> = segments_to_merge
            .iter()
            .map(|s| s.metadata.id.clone())
            .collect();
        sealed.insert(new_id.clone(), new_segment);

        // Update stats
        self.stats
            .compactions_completed
            .fetch_add(1, Ordering::Relaxed);
        self.stats
            .bytes_reclaimed
            .fetch_add(bytes_reclaimed as u64, Ordering::Relaxed);
        self.stats
            .tombstones_collected
            .fetch_add(tombstones_removed as u64, Ordering::Relaxed);

        Some(CompactionResult {
            merged_segments: merged_ids,
            new_segment: Some(new_id),
            vectors_compacted,
            tombstones_removed,
            bytes_reclaimed,
            duration_ms: start.elapsed().as_millis() as u64,
        })
    }

    /// Run automatic compaction based on configuration
    pub fn auto_compact(&self) -> Vec<CompactionResult> {
        let mut results = Vec::new();

        // First, compact segments with high garbage ratio
        let garbage_segments = self.segments_needing_compaction();
        for segment in garbage_segments {
            if let Some(result) = self.compact(&[segment.id]) {
                results.push(result);
            }
        }

        // Then, merge small segments
        let small_segs = self.small_segments();
        if small_segs.len() >= 2 {
            let ids: Vec<_> = small_segs
                .into_iter()
                .take(self.config.max_merge_segments)
                .map(|s| s.id)
                .collect();

            if let Some(result) = self.compact(&ids) {
                results.push(result);
            }
        }

        results
    }

    /// Get statistics
    pub fn stats(&self) -> SegmentStats {
        let active_count = if self.active_segment.read().is_some() {
            1
        } else {
            0
        };
        let sealed = self.sealed_segments.read();

        let total_live: usize = sealed.values().map(|s| s.metadata.live_count).sum();
        let total_deleted: usize = sealed.values().map(|s| s.metadata.deleted_count).sum();
        let total_size: usize = sealed.values().map(|s| s.metadata.size_bytes).sum();

        let active_live = self
            .active_segment
            .read()
            .as_ref()
            .map(|s| s.metadata.live_count)
            .unwrap_or(0);
        let active_size = self
            .active_segment
            .read()
            .as_ref()
            .map(|s| s.metadata.size_bytes)
            .unwrap_or(0);

        SegmentStats {
            namespace: self.namespace.clone(),
            active_segments: active_count,
            sealed_segments: sealed.len(),
            total_live_vectors: total_live + active_live,
            total_deleted_vectors: total_deleted,
            total_size_bytes: total_size + active_size,
            average_segment_size: if !sealed.is_empty() {
                total_live / sealed.len()
            } else {
                0
            },
            garbage_ratio: if total_live + total_deleted > 0 {
                total_deleted as f32 / (total_live + total_deleted) as f32
            } else {
                0.0
            },
        }
    }

    /// Clear all segments
    pub fn clear(&self) {
        *self.active_segment.write() = None;
        self.sealed_segments.write().clear();
        self.vector_index.write().clear();
    }
}

// ============================================================================
// Compaction Manager
// ============================================================================

/// Global compaction manager for all namespaces
pub struct CompactionManager {
    /// Configuration
    config: CompactionConfig,
    /// Per-namespace segment managers
    namespaces: RwLock<HashMap<NamespaceId, Arc<NamespaceSegmentManager>>>,
    /// Global statistics
    global_stats: CompactionStats,
}

impl CompactionManager {
    /// Create a new compaction manager
    pub fn new(config: CompactionConfig) -> Self {
        Self {
            config,
            namespaces: RwLock::new(HashMap::new()),
            global_stats: CompactionStats::default(),
        }
    }

    /// Get or create namespace manager
    pub fn namespace(&self, namespace: &NamespaceId) -> Arc<NamespaceSegmentManager> {
        let mut namespaces = self.namespaces.write();

        if let Some(manager) = namespaces.get(namespace) {
            return Arc::clone(manager);
        }

        let manager = Arc::new(NamespaceSegmentManager::new(
            namespace.clone(),
            self.config.clone(),
        ));
        namespaces.insert(namespace.clone(), Arc::clone(&manager));
        manager
    }

    /// Add a vector to a namespace
    pub fn add(&self, namespace: &NamespaceId, vector: Vector) {
        self.namespace(namespace).add(vector);
    }

    /// Get a vector from a namespace
    pub fn get(&self, namespace: &NamespaceId, id: &VectorId) -> Option<Vector> {
        self.namespaces.read().get(namespace)?.get(id)
    }

    /// Delete a vector from a namespace
    pub fn delete(&self, namespace: &NamespaceId, id: &VectorId) -> bool {
        match self.namespaces.read().get(namespace) {
            Some(manager) => manager.delete(id),
            None => false,
        }
    }

    /// Get all vectors in a namespace
    pub fn get_all(&self, namespace: &NamespaceId) -> Vec<Vector> {
        match self.namespaces.read().get(namespace) {
            Some(manager) => manager.get_all(),
            None => Vec::new(),
        }
    }

    /// Run compaction on a namespace
    pub fn compact_namespace(&self, namespace: &NamespaceId) -> Vec<CompactionResult> {
        match self.namespaces.read().get(namespace) {
            Some(manager) => manager.auto_compact(),
            None => Vec::new(),
        }
    }

    /// Run compaction on all namespaces
    pub fn compact_all(&self) -> HashMap<NamespaceId, Vec<CompactionResult>> {
        let namespaces = self.namespaces.read();
        let mut results = HashMap::new();

        for (namespace, manager) in namespaces.iter() {
            let namespace_results = manager.auto_compact();
            if !namespace_results.is_empty() {
                results.insert(namespace.clone(), namespace_results);
            }
        }

        results
    }

    /// Get statistics for a namespace
    pub fn namespace_stats(&self, namespace: &NamespaceId) -> Option<SegmentStats> {
        self.namespaces.read().get(namespace).map(|m| m.stats())
    }

    /// Get global statistics
    pub fn global_stats(&self) -> GlobalCompactionStats {
        let namespaces = self.namespaces.read();

        let mut total_live = 0usize;
        let mut total_deleted = 0usize;
        let mut total_size = 0usize;
        let mut total_segments = 0usize;

        for manager in namespaces.values() {
            let stats = manager.stats();
            total_live += stats.total_live_vectors;
            total_deleted += stats.total_deleted_vectors;
            total_size += stats.total_size_bytes;
            total_segments += stats.active_segments + stats.sealed_segments;
        }

        GlobalCompactionStats {
            total_namespaces: namespaces.len(),
            total_segments,
            total_live_vectors: total_live,
            total_deleted_vectors: total_deleted,
            total_size_bytes: total_size,
            overall_garbage_ratio: if total_live + total_deleted > 0 {
                total_deleted as f32 / (total_live + total_deleted) as f32
            } else {
                0.0
            },
            compactions_completed: self
                .global_stats
                .compactions_completed
                .load(Ordering::Relaxed),
            bytes_reclaimed: self.global_stats.bytes_reclaimed.load(Ordering::Relaxed),
        }
    }

    /// Delete a namespace
    pub fn delete_namespace(&self, namespace: &NamespaceId) -> bool {
        self.namespaces.write().remove(namespace).is_some()
    }

    /// List all namespaces
    pub fn list_namespaces(&self) -> Vec<NamespaceId> {
        self.namespaces.read().keys().cloned().collect()
    }
}

// ============================================================================
// Statistics Types
// ============================================================================

/// Statistics for compaction operations (atomic)
#[derive(Debug, Default)]
pub struct CompactionStats {
    pub vectors_written: AtomicU64,
    pub vectors_deleted: AtomicU64,
    pub compactions_completed: AtomicU64,
    pub bytes_reclaimed: AtomicU64,
    pub tombstones_collected: AtomicU64,
}

/// Statistics for a namespace's segments
#[derive(Debug, Clone)]
pub struct SegmentStats {
    pub namespace: NamespaceId,
    pub active_segments: usize,
    pub sealed_segments: usize,
    pub total_live_vectors: usize,
    pub total_deleted_vectors: usize,
    pub total_size_bytes: usize,
    pub average_segment_size: usize,
    pub garbage_ratio: f32,
}

/// Global compaction statistics
#[derive(Debug, Clone)]
pub struct GlobalCompactionStats {
    pub total_namespaces: usize,
    pub total_segments: usize,
    pub total_live_vectors: usize,
    pub total_deleted_vectors: usize,
    pub total_size_bytes: usize,
    pub overall_garbage_ratio: f32,
    pub compactions_completed: u64,
    pub bytes_reclaimed: u64,
}

// ============================================================================
// Helper Functions
// ============================================================================

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

/// Estimate size of a vector in bytes
fn estimate_vector_size(vector: &Vector) -> usize {
    let values_size = vector.values.len() * 4;
    let id_size = vector.id.len();
    let metadata_size = vector
        .metadata
        .as_ref()
        .map(|m| serde_json::to_string(m).map(|s| s.len()).unwrap_or(0))
        .unwrap_or(0);

    values_size + id_size + metadata_size + 32 // overhead for struct fields
}

// ============================================================================
// Tests
// ============================================================================

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

    fn make_vector(id: &str, dim: usize) -> Vector {
        Vector {
            id: id.to_string(),
            values: (0..dim).map(|i| i as f32).collect(),
            metadata: None,
            ttl_seconds: None,
            expires_at: None,
        }
    }

    #[test]
    fn test_compaction_config_builder() {
        let config = CompactionConfig::new()
            .with_target_size(5000)
            .with_min_size(500)
            .with_garbage_threshold(0.4)
            .without_auto_compact();

        assert_eq!(config.target_segment_size, 5000);
        assert_eq!(config.min_segment_size, 500);
        assert!((config.garbage_threshold - 0.4).abs() < 0.001);
        assert!(!config.auto_compact);
    }

    #[test]
    fn test_segment_metadata() {
        let mut meta = SegmentMetadata::new("seg1".to_string(), "ns1".to_string());

        assert_eq!(meta.state, SegmentState::Active);
        assert_eq!(meta.live_count, 0);
        assert!(meta.is_empty());
        assert!((meta.garbage_ratio() - 0.0).abs() < 0.001);

        meta.live_count = 100;
        meta.deleted_count = 50;

        assert!(!meta.is_empty());
        assert_eq!(meta.total_count(), 150);
        assert!((meta.garbage_ratio() - 0.333).abs() < 0.01);
        assert!(meta.needs_compaction(0.3));
        assert!(!meta.needs_compaction(0.5));
    }

    #[test]
    fn test_segment_operations() {
        let mut segment = Segment::new("seg1".to_string(), "ns1".to_string());

        // Add vectors
        segment.add(make_vector("v1", 128));
        segment.add(make_vector("v2", 128));
        segment.add(make_vector("v3", 128));

        assert_eq!(segment.metadata.live_count, 3);
        assert!(segment.contains(&"v1".to_string()));
        assert!(segment.get(&"v1".to_string()).is_some());

        // Delete a vector
        assert!(segment.delete(&"v2".to_string()));
        assert_eq!(segment.metadata.live_count, 2);
        assert_eq!(segment.metadata.deleted_count, 1);
        assert!(!segment.contains(&"v2".to_string()));
        assert!(segment.is_tombstoned(&"v2".to_string()));

        // Live vectors
        let live: Vec<_> = segment.live_vectors().collect();
        assert_eq!(live.len(), 2);
    }

    #[test]
    fn test_segment_seal() {
        let mut segment = Segment::new("seg1".to_string(), "ns1".to_string());
        segment.add(make_vector("v1", 128));

        assert_eq!(segment.metadata.state, SegmentState::Active);

        segment.seal();
        assert_eq!(segment.metadata.state, SegmentState::Sealed);
    }

    #[test]
    fn test_namespace_segment_manager() {
        let config = CompactionConfig::new().with_target_size(100);
        let manager = NamespaceSegmentManager::new("ns1".to_string(), config);

        // Add vectors
        for i in 0..50 {
            manager.add(make_vector(&format!("v{}", i), 128));
        }

        assert_eq!(manager.get_all().len(), 50);
        assert!(manager.get(&"v10".to_string()).is_some());

        // Delete some
        assert!(manager.delete(&"v5".to_string()));
        assert!(manager.delete(&"v15".to_string()));

        assert_eq!(manager.get_all().len(), 48);
        assert!(manager.get(&"v5".to_string()).is_none());
    }

    #[test]
    fn test_namespace_manager_auto_seal() {
        let config = CompactionConfig::new()
            .with_target_size(10)
            .with_min_size(5);
        let manager = NamespaceSegmentManager::new("ns1".to_string(), config);

        // Add vectors to trigger segment seal
        for i in 0..25 {
            manager.add(make_vector(&format!("v{}", i), 64));
        }

        let stats = manager.stats();
        // Should have created sealed segments
        assert!(stats.sealed_segments >= 1);
        assert_eq!(stats.total_live_vectors, 25);
    }

    #[test]
    fn test_compaction() {
        let config = CompactionConfig::new()
            .with_target_size(10)
            .with_min_size(5)
            .with_garbage_threshold(0.2);
        let manager = NamespaceSegmentManager::new("ns1".to_string(), config);

        // Add vectors
        for i in 0..30 {
            manager.add(make_vector(&format!("v{}", i), 64));
        }

        // Delete some to create garbage
        for i in 0..10 {
            manager.delete(&format!("v{}", i));
        }

        // Run compaction
        let _results = manager.auto_compact();

        // Verify all remaining vectors are accessible
        assert_eq!(manager.get_all().len(), 20);

        // Check stats
        let stats = manager.stats();
        assert_eq!(stats.total_live_vectors, 20);
    }

    #[test]
    fn test_compaction_manager() {
        let config = CompactionConfig::new().with_target_size(50);
        let manager = CompactionManager::new(config);

        // Add vectors to multiple namespaces
        for i in 0..20 {
            manager.add(&"ns1".to_string(), make_vector(&format!("v{}", i), 128));
            manager.add(&"ns2".to_string(), make_vector(&format!("v{}", i), 128));
        }

        assert_eq!(manager.get_all(&"ns1".to_string()).len(), 20);
        assert_eq!(manager.get_all(&"ns2".to_string()).len(), 20);
        assert_eq!(manager.list_namespaces().len(), 2);

        // Get vector
        assert!(manager.get(&"ns1".to_string(), &"v5".to_string()).is_some());

        // Delete
        assert!(manager.delete(&"ns1".to_string(), &"v5".to_string()));
        assert!(manager.get(&"ns1".to_string(), &"v5".to_string()).is_none());

        // Stats
        let global = manager.global_stats();
        assert_eq!(global.total_namespaces, 2);
        assert_eq!(global.total_live_vectors, 39); // 40 - 1 deleted
    }

    #[test]
    fn test_compaction_result() {
        let config = CompactionConfig::new()
            .with_target_size(5)
            .with_min_size(2)
            .with_garbage_threshold(0.1);
        let manager = NamespaceSegmentManager::new("ns1".to_string(), config);

        // Create segments with garbage
        for i in 0..15 {
            manager.add(make_vector(&format!("v{}", i), 32));
        }

        // Delete to create garbage
        for i in 0..5 {
            manager.delete(&format!("v{}", i));
        }

        // Run compaction
        let _results = manager.auto_compact();

        // Verify data integrity
        assert_eq!(manager.get_all().len(), 10);
        for i in 5..15 {
            assert!(manager.get(&format!("v{}", i)).is_some());
        }
    }

    #[test]
    fn test_segment_stats() {
        let config = CompactionConfig::new().with_target_size(20);
        let manager = NamespaceSegmentManager::new("ns1".to_string(), config);

        for i in 0..30 {
            manager.add(make_vector(&format!("v{}", i), 64));
        }

        // Delete some
        for i in 0..5 {
            manager.delete(&format!("v{}", i));
        }

        let stats = manager.stats();
        assert_eq!(stats.namespace, "ns1");
        assert_eq!(stats.total_live_vectors, 25);
        assert_eq!(stats.total_deleted_vectors, 5);
        assert!(stats.garbage_ratio > 0.0);
    }

    #[test]
    fn test_delete_namespace() {
        let config = CompactionConfig::new();
        let manager = CompactionManager::new(config);

        manager.add(&"ns1".to_string(), make_vector("v1", 128));
        manager.add(&"ns2".to_string(), make_vector("v1", 128));

        assert_eq!(manager.list_namespaces().len(), 2);

        assert!(manager.delete_namespace(&"ns1".to_string()));
        assert_eq!(manager.list_namespaces().len(), 1);

        assert!(!manager.delete_namespace(&"ns1".to_string())); // Already deleted
    }

    #[test]
    fn test_small_segment_merge() {
        let config = CompactionConfig::new()
            .with_target_size(100)
            .with_min_size(20);
        let manager = NamespaceSegmentManager::new("ns1".to_string(), config);

        // Create small segments by adding few vectors then forcing seal
        for i in 0..10 {
            manager.add(make_vector(&format!("v{}", i), 64));
        }

        // Check for small segments
        let _small = manager.small_segments();
        // Might be empty since active segment isn't sealed yet

        // The manager should still work correctly
        assert_eq!(manager.get_all().len(), 10);
    }

    #[test]
    fn test_vector_update() {
        let config = CompactionConfig::new();
        let manager = NamespaceSegmentManager::new("ns1".to_string(), config);

        // Add initial vector
        let v1 = Vector {
            id: "v1".to_string(),
            values: vec![1.0, 2.0, 3.0],
            metadata: None,
            ttl_seconds: None,
            expires_at: None,
        };
        manager.add(v1);

        // Update with new values
        let v1_updated = Vector {
            id: "v1".to_string(),
            values: vec![4.0, 5.0, 6.0],
            metadata: None,
            ttl_seconds: None,
            expires_at: None,
        };
        manager.add(v1_updated);

        // Should still have only one vector
        assert_eq!(manager.get_all().len(), 1);

        // Should have updated values
        let retrieved = manager.get(&"v1".to_string()).unwrap();
        assert_eq!(retrieved.values, vec![4.0, 5.0, 6.0]);
    }
}