ryo_storage/storage/

store.rs

//! Global RYO storage management.
//!
//! Manages the `~/.ryo/` directory structure:
//! ```text
//! ~/.ryo/
//! ├── sessions/           # Session transaction logs
//! │   ├── {id}.txlog.json
//! │   └── ...
//! ├── index.json          # Session metadata index
//! └── config.toml         # (future) Global configuration
//! ```

use super::format::{get_serializer, Format, FormatError};
use super::index::{SessionIndex, SessionMeta};
use super::project::{ProjectIndex, ProjectMeta};
use crate::txlog::TxLog;
use std::fs::{self, File};
use std::io::BufReader;
use std::path::{Path, PathBuf};
use thiserror::Error;

/// Errors from storage operations.
#[derive(Debug, Error)]
pub enum StorageError {
    /// Underlying filesystem I/O failure.
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),

    /// Session payload (de)serialization failure; wraps [`FormatError`].
    #[error("Format error: {0}")]
    Format(#[from] FormatError),

    /// The requested session id does not exist in the store.
    #[error("Session not found: {0}")]
    SessionNotFound(String),

    /// The requested project id does not exist in the store.
    #[error("Project not found: {0}")]
    ProjectNotFound(String),

    /// The storage directory has not been initialized yet (no metadata
    /// file). Carries the directory that was inspected.
    #[error("Storage not initialized at {0}")]
    NotInitialized(PathBuf),

    /// Session/project index lookup or update failure. Carries a rendered
    /// message for diagnostics.
    #[error("Index error: {0}")]
    Index(String),
}

/// Result type for storage operations.
pub type StorageResult<T> = Result<T, StorageError>;

/// Global RYO storage manager.
///
/// Provides high-level API for:
/// - Dumping session logs
/// - Loading past sessions
/// - Querying session history
/// - Managing imported projects
/// - Managing storage lifecycle
#[derive(Debug)]
pub struct RyoStorage {
    /// Root directory (~/.ryo/)
    root: PathBuf,
    /// Sessions directory (~/.ryo/sessions/)
    sessions_dir: PathBuf,
    /// Projects directory (~/.ryo/projects/)
    projects_dir: PathBuf,
    /// Session index file path (~/.ryo/index.json)
    index_path: PathBuf,
    /// Project index file path (~/.ryo/projects.json)
    project_index_path: PathBuf,
    /// Serialization format
    format: Format,
    /// Cached session index (lazy-loaded)
    index: Option<SessionIndex>,
    /// Cached project index (lazy-loaded)
    project_index: Option<ProjectIndex>,
}

impl RyoStorage {
    /// Default RYO directory name.
    pub const DIR_NAME: &'static str = ".ryo";
    /// Sessions subdirectory name.
    pub const SESSIONS_DIR: &'static str = "sessions";
    /// Projects subdirectory name.
    pub const PROJECTS_DIR: &'static str = "projects";
    /// Session index file name.
    pub const INDEX_FILE: &'static str = "index.json";
    /// Project index file name.
    pub const PROJECT_INDEX_FILE: &'static str = "projects.json";

    /// Create a new storage manager at the default location (~/.ryo/).
    pub fn global() -> StorageResult<Self> {
        let home = dirs::home_dir().ok_or_else(|| {
            StorageError::Io(std::io::Error::new(
                std::io::ErrorKind::NotFound,
                "Could not find home directory",
            ))
        })?;
        Self::new(home.join(Self::DIR_NAME))
    }

    /// Create a new storage manager at a custom location.
    pub fn new(root: PathBuf) -> StorageResult<Self> {
        let sessions_dir = root.join(Self::SESSIONS_DIR);
        let projects_dir = root.join(Self::PROJECTS_DIR);
        let index_path = root.join(Self::INDEX_FILE);
        let project_index_path = root.join(Self::PROJECT_INDEX_FILE);

        Ok(Self {
            root,
            sessions_dir,
            projects_dir,
            index_path,
            project_index_path,
            format: Format::default(),
            index: None,
            project_index: None,
        })
    }

    /// Set the serialization format.
    pub fn with_format(mut self, format: Format) -> Self {
        self.format = format;
        self
    }

    /// Get the root directory path.
    pub fn root(&self) -> &Path {
        &self.root
    }

    /// Check if storage is initialized (directories exist).
    pub fn is_initialized(&self) -> bool {
        self.root.exists() && self.sessions_dir.exists()
    }

    /// Initialize storage directories.
    pub fn init(&self) -> StorageResult<()> {
        fs::create_dir_all(&self.sessions_dir)?;
        fs::create_dir_all(&self.projects_dir)?;

        // Create empty session index if it doesn't exist
        if !self.index_path.exists() {
            let index = SessionIndex::new();
            let json = serde_json::to_string_pretty(&index)
                .map_err(|e| StorageError::Format(FormatError::Json(e)))?;
            fs::write(&self.index_path, json)?;
        }

        // Create empty project index if it doesn't exist
        if !self.project_index_path.exists() {
            let index = ProjectIndex::new();
            let json = serde_json::to_string_pretty(&index)
                .map_err(|e| StorageError::Format(FormatError::Json(e)))?;
            fs::write(&self.project_index_path, json)?;
        }

        Ok(())
    }

    /// Ensure storage is initialized, creating directories if needed.
    pub fn ensure_init(&self) -> StorageResult<()> {
        if !self.is_initialized() {
            self.init()?;
        }
        Ok(())
    }

    // ========================================================================
    // Session Dump/Load
    // ========================================================================

    /// Dump a session log to storage.
    ///
    /// Returns the session ID.
    pub fn dump(&mut self, log: &TxLog) -> StorageResult<String> {
        self.ensure_init()?;

        let session_id = log.session_id.clone();
        let filename = self.session_filename(&session_id);
        let path = self.sessions_dir.join(&filename);

        // Serialize and write
        let serializer = get_serializer(self.format);
        let file = File::create(&path)?;
        serializer.serialize_to_file(log, file)?;

        // Update index
        self.add_to_index(log)?;

        Ok(session_id)
    }

    /// Load a session log by ID.
    ///
    /// This method automatically detects the file format by trying all known formats.
    pub fn load(&self, session_id: &str) -> StorageResult<TxLog> {
        // Try the configured format first, then fall back to other formats
        let formats_to_try = [self.format, Format::Json, Format::JsonCompact];

        for format in formats_to_try {
            let filename = format!("{}.txlog.{}", session_id, format.extension());
            let path = self.sessions_dir.join(&filename);

            if path.exists() {
                let file = File::open(&path)?;
                let reader = BufReader::new(file);
                let serializer = get_serializer(format);
                let log = serializer.deserialize_from_reader(reader)?;
                return Ok(log);
            }
        }

        Err(StorageError::SessionNotFound(session_id.to_string()))
    }

    /// Check if a session exists (in any format).
    pub fn exists(&self, session_id: &str) -> bool {
        self.find_session_path(session_id).is_some()
    }

    /// Delete a session (in any format).
    pub fn delete(&mut self, session_id: &str) -> StorageResult<()> {
        if let Some(path) = self.find_session_path(session_id) {
            fs::remove_file(&path)?;
        }

        // Update index
        self.remove_from_index(session_id)?;

        Ok(())
    }

    /// Find the path to a session file, trying all known formats.
    fn find_session_path(&self, session_id: &str) -> Option<PathBuf> {
        for format in [Format::Json, Format::JsonCompact] {
            let filename = format!("{}.txlog.{}", session_id, format.extension());
            let path = self.sessions_dir.join(&filename);
            if path.exists() {
                return Some(path);
            }
        }
        None
    }

    /// Generate filename for a session.
    fn session_filename(&self, session_id: &str) -> String {
        format!("{}.txlog.{}", session_id, self.format.extension())
    }

    // ========================================================================
    // Index Management
    // ========================================================================

    /// Get the session index.
    pub fn index(&mut self) -> StorageResult<&SessionIndex> {
        if self.index.is_none() {
            self.load_index()?;
        }
        Ok(self
            .index
            .as_ref()
            .expect("load_index() above sets self.index to Some"))
    }

    /// Get mutable access to the session index.
    fn index_mut(&mut self) -> StorageResult<&mut SessionIndex> {
        if self.index.is_none() {
            self.load_index()?;
        }
        Ok(self
            .index
            .as_mut()
            .expect("load_index() above sets self.index to Some"))
    }

    /// Load the index from disk.
    fn load_index(&mut self) -> StorageResult<()> {
        if self.index_path.exists() {
            let content = fs::read_to_string(&self.index_path)?;
            let index: SessionIndex = serde_json::from_str(&content)
                .map_err(|e| StorageError::Format(FormatError::Json(e)))?;
            self.index = Some(index);
        } else {
            self.index = Some(SessionIndex::new());
        }
        Ok(())
    }

    /// Save the index to disk.
    fn save_index(&self) -> StorageResult<()> {
        if let Some(ref index) = self.index {
            let json = serde_json::to_string_pretty(index)
                .map_err(|e| StorageError::Format(FormatError::Json(e)))?;
            fs::write(&self.index_path, json)?;
        }
        Ok(())
    }

    /// Add a session to the index.
    fn add_to_index(&mut self, log: &TxLog) -> StorageResult<()> {
        let meta = SessionMeta::from_log(log);
        self.index_mut()?.add(meta);
        self.save_index()?;
        Ok(())
    }

    /// Remove a session from the index.
    fn remove_from_index(&mut self, session_id: &str) -> StorageResult<()> {
        self.index_mut()?.remove(session_id);
        self.save_index()?;
        Ok(())
    }

    // ========================================================================
    // Query API
    // ========================================================================

    /// List all sessions.
    pub fn list_sessions(&mut self) -> StorageResult<Vec<&SessionMeta>> {
        Ok(self.index()?.list())
    }

    /// Find sessions by project path.
    pub fn sessions_for_project(
        &mut self,
        project_path: &Path,
    ) -> StorageResult<Vec<&SessionMeta>> {
        Ok(self.index()?.by_project(project_path))
    }

    /// Get the most recent session.
    pub fn latest_session(&mut self) -> StorageResult<Option<&SessionMeta>> {
        Ok(self.index()?.latest())
    }

    /// Get the most recent session for a project.
    pub fn latest_for_project(
        &mut self,
        project_path: &Path,
    ) -> StorageResult<Option<&SessionMeta>> {
        Ok(self.index()?.latest_for_project(project_path))
    }

    // ========================================================================
    // Maintenance
    // ========================================================================

    /// Clean up old sessions (keep last N per project).
    pub fn cleanup(&mut self, keep_per_project: usize) -> StorageResult<usize> {
        let to_delete = self.index_mut()?.cleanup(keep_per_project);
        let count = to_delete.len();

        for session_id in to_delete {
            let filename = self.session_filename(&session_id);
            let path = self.sessions_dir.join(&filename);
            if path.exists() {
                fs::remove_file(&path)?;
            }
        }

        self.save_index()?;
        Ok(count)
    }

    /// Get total storage size in bytes.
    pub fn storage_size(&self) -> StorageResult<u64> {
        let mut total = 0u64;

        if self.sessions_dir.exists() {
            for entry in fs::read_dir(&self.sessions_dir)? {
                let entry = entry?;
                if entry.file_type()?.is_file() {
                    total += entry.metadata()?.len();
                }
            }
        }

        if self.index_path.exists() {
            total += fs::metadata(&self.index_path)?.len();
        }

        Ok(total)
    }

    // ========================================================================
    // Project Management
    // ========================================================================

    /// Get the project index.
    pub fn project_index(&mut self) -> StorageResult<&ProjectIndex> {
        if self.project_index.is_none() {
            self.load_project_index()?;
        }
        Ok(self
            .project_index
            .as_ref()
            .expect("load_project_index() above sets self.project_index to Some"))
    }

    /// Get mutable access to the project index.
    fn project_index_mut(&mut self) -> StorageResult<&mut ProjectIndex> {
        if self.project_index.is_none() {
            self.load_project_index()?;
        }
        Ok(self
            .project_index
            .as_mut()
            .expect("load_project_index() above sets self.project_index to Some"))
    }

    /// Load the project index from disk.
    fn load_project_index(&mut self) -> StorageResult<()> {
        if self.project_index_path.exists() {
            let content = fs::read_to_string(&self.project_index_path)?;
            let index: ProjectIndex = serde_json::from_str(&content)
                .map_err(|e| StorageError::Format(FormatError::Json(e)))?;
            self.project_index = Some(index);
        } else {
            self.project_index = Some(ProjectIndex::new());
        }
        Ok(())
    }

    /// Save the project index to disk.
    fn save_project_index(&self) -> StorageResult<()> {
        if let Some(ref index) = self.project_index {
            let json = serde_json::to_string_pretty(index)
                .map_err(|e| StorageError::Format(FormatError::Json(e)))?;
            fs::write(&self.project_index_path, json)?;
        }
        Ok(())
    }

    /// Register a project in the index.
    pub fn register_project(&mut self, meta: ProjectMeta) -> StorageResult<String> {
        self.ensure_init()?;
        let project_id = meta.project_id.clone();
        self.project_index_mut()?.add(meta);
        self.save_project_index()?;
        Ok(project_id)
    }

    /// Unregister a project from the index.
    pub fn unregister_project(&mut self, project_id: &str) -> StorageResult<Option<ProjectMeta>> {
        let meta = self.project_index_mut()?.remove(project_id);
        self.save_project_index()?;
        Ok(meta)
    }

    /// Get a project by ID.
    pub fn get_project(&mut self, project_id: &str) -> StorageResult<Option<&ProjectMeta>> {
        Ok(self.project_index()?.get(project_id))
    }

    /// Get a project by path.
    pub fn get_project_by_path(&mut self, path: &Path) -> StorageResult<Option<&ProjectMeta>> {
        Ok(self.project_index()?.get_by_path(path))
    }

    /// Get a mutable project by path.
    pub fn get_project_by_path_mut(
        &mut self,
        path: &Path,
    ) -> StorageResult<Option<&mut ProjectMeta>> {
        // First find the project_id
        let project_id = self
            .project_index()?
            .get_by_path(path)
            .map(|p| p.project_id.clone());

        // Then get mutable reference
        if let Some(id) = project_id {
            Ok(self.project_index_mut()?.get_mut(&id))
        } else {
            Ok(None)
        }
    }

    /// Save the project index (public for external updates).
    pub fn save_projects(&self) -> StorageResult<()> {
        self.save_project_index()
    }

    /// Search projects by name pattern.
    pub fn search_projects_by_name(&mut self, pattern: &str) -> StorageResult<Vec<&ProjectMeta>> {
        Ok(self.project_index()?.search_by_name(pattern))
    }

    /// Check if a project path is already registered.
    pub fn is_project_registered(&mut self, path: &Path) -> StorageResult<bool> {
        Ok(self.project_index()?.contains_path(path))
    }

    /// List all registered projects.
    pub fn list_projects(&mut self) -> StorageResult<Vec<&ProjectMeta>> {
        Ok(self.project_index()?.list())
    }

    /// Touch a project (update last accessed time).
    pub fn touch_project(&mut self, project_id: &str) -> StorageResult<()> {
        if let Some(meta) = self.project_index_mut()?.get_mut(project_id) {
            meta.touch();
            self.save_project_index()?;
        }
        Ok(())
    }

    /// Get project statistics.
    pub fn project_stats(&mut self) -> StorageResult<(usize, usize, usize)> {
        let index = self.project_index()?;
        Ok((index.count(), index.total_files(), index.total_lines()))
    }

    /// Cleanup dead server PIDs from all projects.
    ///
    /// Checks each project's server_pid and clears it if the process is no longer running.
    /// Returns the number of projects that were cleaned up.
    pub fn cleanup_dead_servers(&mut self) -> StorageResult<usize> {
        let cleaned = self.project_index_mut()?.cleanup_dead_servers();
        if cleaned > 0 {
            self.save_project_index()?;
        }
        Ok(cleaned)
    }

    /// List all registered projects, with optional dead server cleanup.
    ///
    /// If `cleanup` is true, dead server PIDs are automatically cleared before listing.
    pub fn list_projects_with_cleanup(
        &mut self,
        cleanup: bool,
    ) -> StorageResult<Vec<&ProjectMeta>> {
        if cleanup {
            self.cleanup_dead_servers()?;
        }
        Ok(self.project_index()?.list())
    }

    // ========================================================================
    // CodeGraph Cache
    // ========================================================================

    /// Cache subdirectory name.
    pub const CACHE_DIR: &'static str = "cache";

    /// Get the cache directory path.
    pub fn cache_dir(&self) -> PathBuf {
        self.root.join(Self::CACHE_DIR)
    }

    /// Ensure cache directory exists.
    fn ensure_cache_dir(&self) -> StorageResult<()> {
        let dir = self.cache_dir();
        if !dir.exists() {
            fs::create_dir_all(&dir)?;
        }
        Ok(())
    }

    /// Save a CodeGraph cache.
    ///
    /// Returns the cache file path.
    pub fn save_graph_cache(&self, project_hash: &str, data: &[u8]) -> StorageResult<PathBuf> {
        self.ensure_cache_dir()?;
        let path = self.cache_dir().join(format!("{}.graph.bin", project_hash));
        fs::write(&path, data)?;
        Ok(path)
    }

    /// Load a CodeGraph cache.
    ///
    /// Returns None if cache doesn't exist.
    pub fn load_graph_cache(&self, project_hash: &str) -> StorageResult<Option<Vec<u8>>> {
        let path = self.cache_dir().join(format!("{}.graph.bin", project_hash));
        if !path.exists() {
            return Ok(None);
        }
        let data = fs::read(&path)?;
        Ok(Some(data))
    }

    /// Delete a CodeGraph cache.
    pub fn delete_graph_cache(&self, project_hash: &str) -> StorageResult<()> {
        let path = self.cache_dir().join(format!("{}.graph.bin", project_hash));
        if path.exists() {
            fs::remove_file(&path)?;
        }
        Ok(())
    }

    /// List all cached project hashes.
    pub fn list_graph_caches(&self) -> StorageResult<Vec<String>> {
        let dir = self.cache_dir();
        if !dir.exists() {
            return Ok(Vec::new());
        }

        let mut hashes = Vec::new();
        for entry in fs::read_dir(&dir)? {
            let entry = entry?;
            let path = entry.path();
            if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
                if name.ends_with(".graph.bin") {
                    let hash = name.trim_end_matches(".graph.bin").to_string();
                    hashes.push(hash);
                }
            }
        }
        Ok(hashes)
    }

    /// Get cache storage size in bytes.
    pub fn cache_size(&self) -> StorageResult<u64> {
        let dir = self.cache_dir();
        if !dir.exists() {
            return Ok(0);
        }

        let mut size = 0u64;
        for entry in fs::read_dir(&dir)? {
            let entry = entry?;
            if let Ok(meta) = entry.metadata() {
                size += meta.len();
            }
        }
        Ok(size)
    }

    /// Clear all graph caches.
    pub fn clear_graph_caches(&self) -> StorageResult<usize> {
        let dir = self.cache_dir();
        if !dir.exists() {
            return Ok(0);
        }

        let mut count = 0;
        for entry in fs::read_dir(&dir)? {
            let entry = entry?;
            let path = entry.path();
            if path.extension().map(|e| e == "bin").unwrap_or(false) {
                fs::remove_file(&path)?;
                count += 1;
            }
        }
        Ok(count)
    }
}

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

    fn create_test_log(project: &str) -> TxLog {
        let mut log = TxLog::with_project(project);
        log.log(TxAction::GoalSet {
            query: "test".to_string(),
            intent_type: "test".to_string(),
            confidence: 0.9,
        });
        log
    }

    #[test]
    fn test_init_and_dump() {
        let temp = TempDir::new().unwrap();
        let mut storage = RyoStorage::new(temp.path().join(".ryo")).unwrap();

        storage.init().unwrap();
        assert!(storage.is_initialized());

        let log = create_test_log("/test/project");
        let session_id = storage.dump(&log).unwrap();

        assert!(storage.exists(&session_id));
    }

    #[test]
    fn test_load_session() {
        let temp = TempDir::new().unwrap();
        let mut storage = RyoStorage::new(temp.path().join(".ryo")).unwrap();
        storage.init().unwrap();

        let log = create_test_log("/test/project");
        let session_id = storage.dump(&log).unwrap();

        let loaded = storage.load(&session_id).unwrap();
        assert_eq!(loaded.session_id, log.session_id);
        assert_eq!(loaded.entries().len(), log.entries().len());
    }

    #[test]
    fn test_session_index() {
        let temp = TempDir::new().unwrap();
        let mut storage = RyoStorage::new(temp.path().join(".ryo")).unwrap();
        storage.init().unwrap();

        // Dump multiple sessions
        let log1 = create_test_log("/project/a");
        let log2 = create_test_log("/project/b");
        let log3 = create_test_log("/project/a");

        storage.dump(&log1).unwrap();
        storage.dump(&log2).unwrap();
        storage.dump(&log3).unwrap();

        // List all
        let all = storage.list_sessions().unwrap();
        assert_eq!(all.len(), 3);

        // Filter by project
        let proj_a = storage
            .sessions_for_project(Path::new("/project/a"))
            .unwrap();
        assert_eq!(proj_a.len(), 2);
    }

    #[test]
    fn test_delete_session() {
        let temp = TempDir::new().unwrap();
        let mut storage = RyoStorage::new(temp.path().join(".ryo")).unwrap();
        storage.init().unwrap();

        let log = create_test_log("/test/project");
        let session_id = storage.dump(&log).unwrap();

        assert!(storage.exists(&session_id));
        storage.delete(&session_id).unwrap();
        assert!(!storage.exists(&session_id));
    }
}