aurora_semantic/storage/

disk.rs

//! Disk-based storage implementation.

use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};

use parking_lot::RwLock;

use crate::error::{Error, Result};
use crate::storage::{Storage, StorageOptions, WorkspaceMetadata};
use crate::types::{Chunk, Document, LanguageStats, WorkspaceId, WorkspaceStats};

/// File-based storage implementation.
pub struct DiskStorage {
    /// Root directory for all storage.
    root_dir: PathBuf,
    /// Storage options.
    options: StorageOptions,
    /// Cached metadata.
    metadata_cache: RwLock<HashMap<WorkspaceId, WorkspaceMetadata>>,
}

impl DiskStorage {
    /// Create a new disk storage instance.
    pub fn new(root_dir: PathBuf) -> Result<Self> {
        Self::with_options(root_dir, StorageOptions::default())
    }

    /// Create a new disk storage with custom options.
    pub fn with_options(root_dir: PathBuf, options: StorageOptions) -> Result<Self> {
        // Create root directory if it doesn't exist
        fs::create_dir_all(&root_dir)?;

        Ok(Self {
            root_dir,
            options,
            metadata_cache: RwLock::new(HashMap::new()),
        })
    }

    /// Get the directory for a specific workspace.
    fn workspace_dir(&self, workspace_id: &WorkspaceId) -> PathBuf {
        self.root_dir.join(workspace_id.to_string())
    }

    /// Get the metadata file path for a workspace.
    fn metadata_path(&self, workspace_id: &WorkspaceId) -> PathBuf {
        self.workspace_dir(workspace_id).join("metadata.json")
    }

    /// Get the documents file path for a workspace.
    fn documents_path(&self, workspace_id: &WorkspaceId) -> PathBuf {
        self.workspace_dir(workspace_id).join("documents.bin")
    }

    /// Get the chunks file path for a workspace.
    fn chunks_path(&self, workspace_id: &WorkspaceId) -> PathBuf {
        self.workspace_dir(workspace_id).join("chunks.bin")
    }

    /// Get the embeddings file path for a workspace.
    fn embeddings_path(&self, workspace_id: &WorkspaceId) -> PathBuf {
        self.workspace_dir(workspace_id).join("embeddings.bin")
    }

    /// Serialize data to a file.
    fn write_bincode<T: serde::Serialize>(&self, path: &Path, data: &T) -> Result<()> {
        let file = File::create(path)?;
        let mut writer = BufWriter::with_capacity(self.options.buffer_size, file);
        bincode::serialize_into(&mut writer, data)?;
        writer.flush()?;
        Ok(())
    }

    /// Deserialize data from a file.
    fn read_bincode<T: serde::de::DeserializeOwned>(&self, path: &Path) -> Result<T> {
        let file = File::open(path)?;
        let reader = BufReader::with_capacity(self.options.buffer_size, file);
        let data = bincode::deserialize_from(reader)?;
        Ok(data)
    }

    /// Write JSON data to a file.
    fn write_json<T: serde::Serialize>(&self, path: &Path, data: &T) -> Result<()> {
        let file = File::create(path)?;
        let writer = BufWriter::with_capacity(self.options.buffer_size, file);
        serde_json::to_writer_pretty(writer, data)?;
        Ok(())
    }

    /// Read JSON data from a file.
    fn read_json<T: serde::de::DeserializeOwned>(&self, path: &Path) -> Result<T> {
        let file = File::open(path)?;
        let reader = BufReader::with_capacity(self.options.buffer_size, file);
        let data = serde_json::from_reader(reader)?;
        Ok(data)
    }

    /// Calculate directory size recursively.
    fn dir_size(&self, path: &Path) -> u64 {
        if !path.exists() {
            return 0;
        }

        walkdir::WalkDir::new(path)
            .into_iter()
            .filter_map(|e| e.ok())
            .filter(|e| e.file_type().is_file())
            .filter_map(|e| e.metadata().ok())
            .map(|m| m.len())
            .sum()
    }
}

impl Storage for DiskStorage {
    fn init_workspace(&self, workspace_id: &WorkspaceId, root_path: &PathBuf) -> Result<()> {
        let ws_dir = self.workspace_dir(workspace_id);
        fs::create_dir_all(&ws_dir)?;

        // Create subdirectories
        fs::create_dir_all(self.tantivy_index_path(workspace_id))?;

        // Create initial metadata
        let metadata = WorkspaceMetadata::new(
            workspace_id.clone(),
            root_path.clone(),
            384, // Default embedding dimension
        );

        self.save_workspace_metadata(workspace_id, &metadata)?;

        Ok(())
    }

    fn workspace_exists(&self, workspace_id: &WorkspaceId) -> bool {
        self.metadata_path(workspace_id).exists()
    }

    fn load_workspace_metadata(&self, workspace_id: &WorkspaceId) -> Result<WorkspaceMetadata> {
        // Check cache first
        {
            let cache = self.metadata_cache.read();
            if let Some(meta) = cache.get(workspace_id) {
                return Ok(meta.clone());
            }
        }

        let path = self.metadata_path(workspace_id);
        if !path.exists() {
            return Err(Error::WorkspaceNotFound(workspace_id.clone()));
        }

        let metadata: WorkspaceMetadata = self.read_json(&path)?;

        // Update cache
        {
            let mut cache = self.metadata_cache.write();
            cache.insert(workspace_id.clone(), metadata.clone());
        }

        Ok(metadata)
    }

    fn save_workspace_metadata(
        &self,
        workspace_id: &WorkspaceId,
        metadata: &WorkspaceMetadata,
    ) -> Result<()> {
        let path = self.metadata_path(workspace_id);
        self.write_json(&path, metadata)?;

        // Update cache
        {
            let mut cache = self.metadata_cache.write();
            cache.insert(workspace_id.clone(), metadata.clone());
        }

        Ok(())
    }

    fn save_documents(&self, workspace_id: &WorkspaceId, documents: &[Document]) -> Result<()> {
        let path = self.documents_path(workspace_id);
        self.write_bincode(&path, &documents.to_vec())
    }

    fn load_documents(&self, workspace_id: &WorkspaceId) -> Result<Vec<Document>> {
        let path = self.documents_path(workspace_id);
        if !path.exists() {
            return Ok(Vec::new());
        }
        self.read_bincode(&path)
    }

    fn save_chunks(&self, workspace_id: &WorkspaceId, chunks: &[Chunk]) -> Result<()> {
        let path = self.chunks_path(workspace_id);
        self.write_bincode(&path, &chunks.to_vec())
    }

    fn load_chunks(&self, workspace_id: &WorkspaceId) -> Result<Vec<Chunk>> {
        let path = self.chunks_path(workspace_id);
        if !path.exists() {
            return Ok(Vec::new());
        }
        self.read_bincode(&path)
    }

    fn save_embeddings(
        &self,
        workspace_id: &WorkspaceId,
        embeddings: &[(String, Vec<f32>)],
    ) -> Result<()> {
        let path = self.embeddings_path(workspace_id);
        self.write_bincode(&path, &embeddings.to_vec())
    }

    fn load_embeddings(&self, workspace_id: &WorkspaceId) -> Result<Vec<(String, Vec<f32>)>> {
        let path = self.embeddings_path(workspace_id);
        if !path.exists() {
            return Ok(Vec::new());
        }
        self.read_bincode(&path)
    }

    fn tantivy_index_path(&self, workspace_id: &WorkspaceId) -> PathBuf {
        self.workspace_dir(workspace_id).join("tantivy")
    }

    fn vector_index_path(&self, workspace_id: &WorkspaceId) -> PathBuf {
        self.workspace_dir(workspace_id).join("vectors.usearch")
    }

    fn delete_workspace(&self, workspace_id: &WorkspaceId) -> Result<()> {
        let ws_dir = self.workspace_dir(workspace_id);
        if ws_dir.exists() {
            fs::remove_dir_all(&ws_dir)?;
        }

        // Remove from cache
        {
            let mut cache = self.metadata_cache.write();
            cache.remove(workspace_id);
        }

        Ok(())
    }

    fn list_workspaces(&self) -> Result<Vec<WorkspaceId>> {
        let mut workspaces = Vec::new();

        if !self.root_dir.exists() {
            return Ok(workspaces);
        }

        for entry in fs::read_dir(&self.root_dir)? {
            let entry = entry?;
            if entry.file_type()?.is_dir() {
                if let Some(name) = entry.file_name().to_str() {
                    if let Ok(uuid) = uuid::Uuid::parse_str(name) {
                        workspaces.push(WorkspaceId::from_uuid(uuid));
                    }
                }
            }
        }

        Ok(workspaces)
    }

    fn get_workspace_stats(&self, workspace_id: &WorkspaceId) -> Result<WorkspaceStats> {
        let metadata = self.load_workspace_metadata(workspace_id)?;

        let index_size_bytes = self.dir_size(&self.workspace_dir(workspace_id));

        let language_stats: Vec<LanguageStats> = metadata
            .language_stats
            .iter()
            .map(|(lang, info)| LanguageStats {
                language: *lang,
                file_count: info.file_count,
                chunk_count: info.chunk_count,
                total_bytes: info.total_bytes,
            })
            .collect();

        Ok(WorkspaceStats {
            workspace_id: workspace_id.clone(),
            root_path: metadata.root_path,
            document_count: metadata.document_count,
            chunk_count: metadata.chunk_count,
            total_bytes: metadata.total_bytes,
            index_size_bytes,
            created_at: metadata.created_at,
            updated_at: metadata.updated_at,
            language_stats,
        })
    }

    fn flush(&self) -> Result<()> {
        // Write all cached metadata to disk
        let cache = self.metadata_cache.read();
        for (workspace_id, metadata) in cache.iter() {
            let path = self.metadata_path(workspace_id);
            self.write_json(&path, metadata)?;
        }
        Ok(())
    }
}

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

    #[test]
    fn test_disk_storage_lifecycle() {
        let temp_dir = TempDir::new().unwrap();
        let storage = DiskStorage::new(temp_dir.path().to_path_buf()).unwrap();

        let workspace_id = WorkspaceId::new();
        let root_path = PathBuf::from("/test/project");

        // Initialize workspace
        storage.init_workspace(&workspace_id, &root_path).unwrap();
        assert!(storage.workspace_exists(&workspace_id));

        // Load metadata
        let metadata = storage.load_workspace_metadata(&workspace_id).unwrap();
        assert_eq!(metadata.root_path, root_path);

        // List workspaces
        let workspaces = storage.list_workspaces().unwrap();
        assert!(workspaces.contains(&workspace_id));

        // Delete workspace
        storage.delete_workspace(&workspace_id).unwrap();
        assert!(!storage.workspace_exists(&workspace_id));
    }

    #[test]
    fn test_document_persistence() {
        let temp_dir = TempDir::new().unwrap();
        let storage = DiskStorage::new(temp_dir.path().to_path_buf()).unwrap();

        let workspace_id = WorkspaceId::new();
        storage
            .init_workspace(&workspace_id, &PathBuf::from("/test"))
            .unwrap();

        let docs = vec![Document {
            id: crate::types::DocumentId::new(),
            relative_path: PathBuf::from("test.rs"),
            absolute_path: PathBuf::from("/test/test.rs"),
            language: Language::Rust,
            content_hash: "abc123".to_string(),
            size_bytes: 1000,
            modified_at: chrono::Utc::now(),
        }];

        storage.save_documents(&workspace_id, &docs).unwrap();

        let loaded = storage.load_documents(&workspace_id).unwrap();
        assert_eq!(loaded.len(), 1);
        assert_eq!(loaded[0].relative_path, PathBuf::from("test.rs"));
    }
}