//! Optuna Study Persistence with MinIO
//!
//! This module provides functionality to save and load Optuna studies (SQLite databases)
//! to/from MinIO for crash recovery and distributed hyperparameter tuning.
//!
//! # Architecture
//!
//! - **Local Storage**: Optuna uses SQLite JournalStorage for fast trial recording
//! - **Remote Backup**: After each trial, SQLite file is uploaded to MinIO
//! - **Crash Recovery**: On startup, download existing study from MinIO if available
//! - **Versioning**: MinIO versioning ensures no data loss during concurrent uploads
//!
//! # Study Naming Convention
//!
//! Studies are named: `tuning_{model_type}_{job_id}`
//! - `model_type`: mamba2, dqn, ppo, tft, liquid
//! - `job_id`: UUID for this tuning job
//!
//! Example: `tuning_mamba2_a3c4e5f7-1234-5678-90ab-cdef12345678`
//!
//! # MinIO Structure
//!
//! ```text
//! foxhunt-ml/
//! ├── studies/
//! │   ├── tuning_mamba2_<job_id>.db
//! │   ├── tuning_dqn_<job_id>.db
//! │   └── ...
//! └── models/  (existing model checkpoints)
//! ```

use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;

use storage::{ObjectStoreBackend, Storage};
use tokio::fs;
use tracing::{debug, error, info, warn};

/// Result type for Optuna persistence operations
pub type PersistenceResult<T> = Result<T, PersistenceError>;

/// Errors that can occur during Optuna study persistence
#[derive(Debug, thiserror::Error)]
pub enum PersistenceError {
    /// Study not found in MinIO (first run)
    #[error("Study not found: {study_name}")]
    StudyNotFound { study_name: String },

    /// Network error during MinIO operations
    #[error("Network error: {message}")]
    NetworkError { message: String },

    /// File I/O error
    #[error("I/O error: {message}")]
    IoError { message: String },

    /// Corrupted study file
    #[error("Corrupted study file: {study_name} - {reason}")]
    CorruptedStudy { study_name: String, reason: String },

    /// Storage backend error
    #[error("Storage error: {0}")]
    StorageError(#[from] storage::StorageError),

    /// Invalid study name format
    #[error("Invalid study name: {name} - must be format 'tuning_{{model}}_{{job_id}}'")]
    InvalidStudyName { name: String },
}

impl PersistenceError {
    /// Check if error is retryable
    pub fn is_retryable(&self) -> bool {
        matches!(
            self,
            PersistenceError::NetworkError { .. } | PersistenceError::StorageError(_)
        )
    }
}

/// Configuration for Optuna study persistence
#[derive(Debug, Clone)]
pub struct OptunaPersistenceConfig {
    /// MinIO bucket name (default: "foxhunt-ml")
    pub bucket_name: String,
    /// Study prefix in bucket (default: "studies/")
    pub study_prefix: String,
    /// Local directory for downloaded studies (default: "/tmp/optuna_studies")
    pub local_cache_dir: PathBuf,
    /// Maximum retry attempts for network operations
    pub max_retries: u32,
    /// Initial retry delay
    pub initial_retry_delay: Duration,
    /// Exponential backoff multiplier
    pub backoff_multiplier: f64,
}

impl Default for OptunaPersistenceConfig {
    fn default() -> Self {
        Self {
            bucket_name: "foxhunt-ml".to_owned(),
            study_prefix: "studies/".to_owned(),
            local_cache_dir: PathBuf::from("/tmp/optuna_studies"),
            max_retries: 3,
            initial_retry_delay: Duration::from_millis(100),
            backoff_multiplier: 2.0,
        }
    }
}

/// Optuna Study Persistence Manager
///
/// Handles saving and loading Optuna studies (SQLite databases) to/from MinIO
/// with retry logic, corruption detection, and crash recovery.
pub struct OptunaPersistence {
    /// MinIO storage backend
    storage: Arc<ObjectStoreBackend>,
    /// Configuration
    config: OptunaPersistenceConfig,
}

impl OptunaPersistence {
    /// Create a new Optuna persistence manager
    ///
    /// # Arguments
    ///
    /// * `storage` - MinIO storage backend (already configured for foxhunt-ml bucket)
    /// * `config` - Persistence configuration
    ///
    /// # Errors
    ///
    /// Returns error if local cache directory cannot be created.
    pub async fn new(
        storage: Arc<ObjectStoreBackend>,
        config: OptunaPersistenceConfig,
    ) -> PersistenceResult<Self> {
        // Create local cache directory if it doesn't exist
        fs::create_dir_all(&config.local_cache_dir)
            .await
            .map_err(|e| PersistenceError::IoError {
                message: format!("Failed to create cache directory: {}", e),
            })?;

        info!(
            "Initialized Optuna persistence (bucket: {}, prefix: {})",
            config.bucket_name, config.study_prefix
        );

        Ok(Self { storage, config })
    }

    /// Save Optuna study to MinIO
    ///
    /// Uploads the SQLite database file to MinIO with automatic retry on failure.
    /// Safe for concurrent calls - MinIO versioning prevents data loss.
    ///
    /// # Arguments
    ///
    /// * `study_name` - Study name (format: `tuning_{model_type}_{job_id}`)
    /// * `local_db_path` - Path to local SQLite database file
    ///
    /// # Errors
    ///
    /// Returns error if:
    /// - Study name format is invalid
    /// - Local database file doesn't exist or can't be read
    /// - Network errors exceed retry limit
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use std::path::PathBuf;
    /// # use std::sync::Arc;
    /// # async fn example(persistence: Arc<OptunaPersistence>) -> Result<(), Box<dyn std::error::Error>> {
    /// let study_name = "tuning_mamba2_a3c4e5f7-1234-5678-90ab-cdef12345678";
    /// let db_path = PathBuf::from("/tmp/optuna.db");
    /// persistence.save_study_to_minio(study_name, &db_path).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn save_study_to_minio(
        &self,
        study_name: &str,
        local_db_path: &Path,
    ) -> PersistenceResult<()> {
        // Validate study name format
        self.validate_study_name(study_name)?;

        // Check if local database exists
        if !local_db_path.exists() {
            return Err(PersistenceError::IoError {
                message: format!("Study database not found: {}", local_db_path.display()),
            });
        }

        // Read SQLite database file
        let db_data = fs::read(local_db_path)
            .await
            .map_err(|e| PersistenceError::IoError {
                message: format!("Failed to read study database: {}", e),
            })?;

        info!(
            "Saving study '{}' to MinIO ({} bytes)",
            study_name,
            db_data.len()
        );

        // Construct MinIO path
        let minio_path = format!("{}{}.db", self.config.study_prefix, study_name);

        // Upload with retry logic
        self.upload_with_retry(&minio_path, &db_data, study_name)
            .await?;

        info!(
            "Successfully saved study '{}' to MinIO (path: {})",
            study_name, minio_path
        );

        Ok(())
    }

    /// Load existing study from MinIO
    ///
    /// Downloads SQLite database from MinIO to local cache directory.
    /// Validates file integrity before returning.
    ///
    /// # Arguments
    ///
    /// * `study_name` - Study name to load
    ///
    /// # Returns
    ///
    /// Returns path to local SQLite database file.
    ///
    /// # Errors
    ///
    /// Returns error if:
    /// - Study name format is invalid
    /// - Study not found in MinIO (first run)
    /// - Download fails after retries
    /// - Downloaded file is corrupted
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use std::sync::Arc;
    /// # async fn example(persistence: Arc<OptunaPersistence>) -> Result<(), Box<dyn std::error::Error>> {
    /// let study_name = "tuning_mamba2_a3c4e5f7-1234-5678-90ab-cdef12345678";
    /// match persistence.load_study_from_minio(study_name).await {
    ///     Ok(db_path) => {
    ///         println!("Loaded study from: {}", db_path.display());
    ///         // Use db_path with Optuna JournalStorage
    ///     }
    ///     Err(e) if matches!(e, optuna_persistence::PersistenceError::StudyNotFound { .. }) => {
    ///         println!("First run - creating new study");
    ///     }
    ///     Err(e) => return Err(e.into()),
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn load_study_from_minio(&self, study_name: &str) -> PersistenceResult<PathBuf> {
        // Validate study name format
        self.validate_study_name(study_name)?;

        info!("Loading study '{}' from MinIO", study_name);

        // Construct MinIO path
        let minio_path = format!("{}{}.db", self.config.study_prefix, study_name);

        // Check if study exists in MinIO
        let exists = self.storage.exists(&minio_path).await?;
        if !exists {
            return Err(PersistenceError::StudyNotFound {
                study_name: study_name.to_owned(),
            });
        }

        // Download with retry logic
        let db_data = self.download_with_retry(&minio_path, study_name).await?;

        // Validate SQLite database format
        self.validate_sqlite_format(&db_data, study_name)?;

        // Write to local cache
        let local_path = self
            .config
            .local_cache_dir
            .join(format!("{}.db", study_name));
        fs::write(&local_path, &db_data)
            .await
            .map_err(|e| PersistenceError::IoError {
                message: format!("Failed to write study to cache: {}", e),
            })?;

        info!(
            "Successfully loaded study '{}' from MinIO ({} bytes) -> {}",
            study_name,
            db_data.len(),
            local_path.display()
        );

        Ok(local_path)
    }

    /// List all available studies in MinIO
    ///
    /// Returns list of study names (without .db extension).
    ///
    /// # Errors
    ///
    /// Returns error if MinIO list operation fails.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use std::sync::Arc;
    /// # async fn example(persistence: Arc<OptunaPersistence>) -> Result<(), Box<dyn std::error::Error>> {
    /// let studies = persistence.list_studies().await?;
    /// for study in studies {
    ///     println!("Found study: {}", study);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn list_studies(&self) -> PersistenceResult<Vec<String>> {
        info!(
            "Listing studies from MinIO (prefix: {})",
            self.config.study_prefix
        );

        // List all objects with study prefix
        let paths = self.storage.list(&self.config.study_prefix).await?;

        // Extract study names (remove prefix and .db extension)
        let studies: Vec<String> = paths
            .iter()
            .filter_map(|path| {
                // Remove prefix
                let name = path.strip_prefix(&self.config.study_prefix)?;
                // Remove .db extension
                name.strip_suffix(".db")
            })
            .map(String::from)
            .collect();

        info!("Found {} studies in MinIO", studies.len());
        debug!("Studies: {:?}", studies);

        Ok(studies)
    }

    /// Delete study from MinIO
    ///
    /// Removes study database from MinIO. Does not affect local cache.
    ///
    /// # Arguments
    ///
    /// * `study_name` - Study name to delete
    ///
    /// # Returns
    ///
    /// Returns `true` if study was deleted, `false` if not found.
    ///
    /// # Errors
    ///
    /// Returns error if MinIO delete operation fails.
    pub async fn delete_study(&self, study_name: &str) -> PersistenceResult<bool> {
        self.validate_study_name(study_name)?;

        info!("Deleting study '{}' from MinIO", study_name);

        let minio_path = format!("{}{}.db", self.config.study_prefix, study_name);
        let deleted = self.storage.delete(&minio_path).await?;

        if deleted {
            info!("Successfully deleted study '{}'", study_name);
        } else {
            warn!("Study '{}' not found for deletion", study_name);
        }

        Ok(deleted)
    }

    // --- Private Helper Methods ---

    /// Validate study name format
    fn validate_study_name(&self, study_name: &str) -> PersistenceResult<()> {
        // Expected format: tuning_{model_type}_{job_id}
        if !study_name.starts_with("tuning_") {
            return Err(PersistenceError::InvalidStudyName {
                name: study_name.to_owned(),
            });
        }

        let parts: Vec<&str> = study_name.split('_').collect();
        if parts.len() < 3 {
            return Err(PersistenceError::InvalidStudyName {
                name: study_name.to_owned(),
            });
        }

        Ok(())
    }

    /// Upload data to MinIO with retry logic
    async fn upload_with_retry(
        &self,
        path: &str,
        data: &[u8],
        study_name: &str,
    ) -> PersistenceResult<()> {
        let mut delay = self.config.initial_retry_delay;

        for attempt in 1..=self.config.max_retries {
            match self.storage.store(path, data).await {
                Ok(_) => {
                    if attempt > 1 {
                        info!(
                            "Successfully uploaded study '{}' after {} attempts",
                            study_name, attempt
                        );
                    }
                    return Ok(());
                },
                Err(e) => {
                    if attempt < self.config.max_retries {
                        warn!(
                            "Upload attempt {}/{} failed for study '{}': {} - retrying in {:?}",
                            attempt, self.config.max_retries, study_name, e, delay
                        );
                        tokio::time::sleep(delay).await;
                        delay = delay.mul_f32(self.config.backoff_multiplier as f32);
                    } else {
                        error!(
                            "Failed to upload study '{}' after {} attempts: {}",
                            study_name, self.config.max_retries, e
                        );
                        return Err(PersistenceError::NetworkError {
                            message: format!("Upload failed after {} attempts: {}", attempt, e),
                        });
                    }
                },
            }
        }

        Err(PersistenceError::NetworkError {
            message: format!(
                "Upload retry loop for study '{}' exhausted without returning",
                study_name
            ),
        })
    }

    /// Download data from MinIO with retry logic
    async fn download_with_retry(
        &self,
        path: &str,
        study_name: &str,
    ) -> PersistenceResult<Vec<u8>> {
        let mut delay = self.config.initial_retry_delay;

        for attempt in 1..=self.config.max_retries {
            match self.storage.retrieve(path).await {
                Ok(data) => {
                    if attempt > 1 {
                        info!(
                            "Successfully downloaded study '{}' after {} attempts",
                            study_name, attempt
                        );
                    }
                    return Ok(data);
                },
                Err(e) => {
                    if attempt < self.config.max_retries {
                        warn!(
                            "Download attempt {}/{} failed for study '{}': {} - retrying in {:?}",
                            attempt, self.config.max_retries, study_name, e, delay
                        );
                        tokio::time::sleep(delay).await;
                        delay = delay.mul_f32(self.config.backoff_multiplier as f32);
                    } else {
                        error!(
                            "Failed to download study '{}' after {} attempts: {}",
                            study_name, self.config.max_retries, e
                        );
                        return Err(PersistenceError::NetworkError {
                            message: format!("Download failed after {} attempts: {}", attempt, e),
                        });
                    }
                },
            }
        }

        Err(PersistenceError::NetworkError {
            message: format!(
                "Download retry loop for study '{}' exhausted without returning",
                study_name
            ),
        })
    }

    /// Validate SQLite database format
    ///
    /// Basic validation to detect corrupted files before writing to disk.
    fn validate_sqlite_format(&self, data: &[u8], study_name: &str) -> PersistenceResult<()> {
        // SQLite files start with magic header: "SQLite format 3\0"
        const SQLITE_MAGIC: &[u8] = b"SQLite format 3\0";

        if data.len() < SQLITE_MAGIC.len() {
            return Err(PersistenceError::CorruptedStudy {
                study_name: study_name.to_owned(),
                reason: format!("File too small ({} bytes)", data.len()),
            });
        }

        if &data[..SQLITE_MAGIC.len()] != SQLITE_MAGIC {
            return Err(PersistenceError::CorruptedStudy {
                study_name: study_name.to_owned(),
                reason: "Invalid SQLite header".to_owned(),
            });
        }

        debug!("SQLite format validation passed for study '{}'", study_name);
        Ok(())
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use tempfile::NamedTempFile;

    // Helper to create test SQLite database
    async fn create_test_db() -> NamedTempFile {
        let temp_file = NamedTempFile::new().expect("Failed to create temp file");
        let db_path = temp_file.path();

        // Write SQLite magic header
        fs::write(db_path, b"SQLite format 3\0test data")
            .await
            .expect("Failed to write test database");

        temp_file
    }

    #[test]
    fn test_validate_study_name() {
        let config = OptunaPersistenceConfig::default();
        let storage = Arc::new(
            storage::object_store_backend::test_helpers::new_for_testing(
                Arc::new(object_store::memory::InMemory::new()),
                "foxhunt-ml".to_owned(),
            ),
        );
        let persistence = OptunaPersistence { storage, config };

        // Valid names
        assert!(persistence
            .validate_study_name("tuning_mamba2_a3c4e5f7-1234-5678-90ab-cdef12345678")
            .is_ok());
        assert!(persistence.validate_study_name("tuning_dqn_12345").is_ok());

        // Invalid names
        assert!(persistence.validate_study_name("invalid").is_err());
        assert!(persistence.validate_study_name("tuning_only").is_err());
        assert!(persistence.validate_study_name("mamba2_job123").is_err());
    }

    #[test]
    fn test_validate_sqlite_format() {
        let config = OptunaPersistenceConfig::default();
        let storage = Arc::new(
            storage::object_store_backend::test_helpers::new_for_testing(
                Arc::new(object_store::memory::InMemory::new()),
                "foxhunt-ml".to_owned(),
            ),
        );
        let persistence = OptunaPersistence { storage, config };

        // Valid SQLite header
        let valid_data = b"SQLite format 3\0some data";
        assert!(persistence
            .validate_sqlite_format(valid_data, "test_study")
            .is_ok());

        // Invalid header
        let invalid_data = b"Not SQLite";
        assert!(persistence
            .validate_sqlite_format(invalid_data, "test_study")
            .is_err());

        // Too short
        let short_data = b"SQLite";
        assert!(persistence
            .validate_sqlite_format(short_data, "test_study")
            .is_err());
    }

    #[tokio::test]
    async fn test_save_and_load_study() {
        // Create in-memory storage backend for testing
        let memory_store = Arc::new(object_store::memory::InMemory::new());
        let storage = Arc::new(
            storage::object_store_backend::test_helpers::new_for_testing(
                memory_store,
                "foxhunt-ml".to_owned(),
            ),
        );

        let config = OptunaPersistenceConfig {
            local_cache_dir: PathBuf::from("/tmp/test_optuna"),
            ..Default::default()
        };

        let persistence = OptunaPersistence::new(storage, config)
            .await
            .expect("Failed to create persistence");

        // Create test database
        let temp_db = create_test_db().await;
        let study_name = "tuning_mamba2_test123";

        // Save study
        persistence
            .save_study_to_minio(study_name, temp_db.path())
            .await
            .expect("Failed to save study");

        // Load study
        let loaded_path = persistence
            .load_study_from_minio(study_name)
            .await
            .expect("Failed to load study");

        // Verify loaded file exists and has correct content
        assert!(loaded_path.exists());
        let loaded_data = fs::read(&loaded_path)
            .await
            .expect("Failed to read loaded file");
        let original_data = fs::read(temp_db.path())
            .await
            .expect("Failed to read original file");
        assert_eq!(loaded_data, original_data);
    }

    #[tokio::test]
    async fn test_list_studies() {
        // Create in-memory storage backend
        let memory_store = Arc::new(object_store::memory::InMemory::new());
        let storage = Arc::new(
            storage::object_store_backend::test_helpers::new_for_testing(
                memory_store,
                "foxhunt-ml".to_owned(),
            ),
        );

        let config = OptunaPersistenceConfig::default();
        let persistence = OptunaPersistence::new(storage.clone(), config)
            .await
            .expect("Failed to create persistence");

        // Save multiple test studies
        let temp_db = create_test_db().await;
        let studies = vec![
            "tuning_mamba2_test1",
            "tuning_dqn_test2",
            "tuning_ppo_test3",
        ];

        for study in &studies {
            persistence
                .save_study_to_minio(study, temp_db.path())
                .await
                .expect("Failed to save study");
        }

        // List studies
        let listed = persistence
            .list_studies()
            .await
            .expect("Failed to list studies");

        // Verify all studies are listed
        assert_eq!(listed.len(), studies.len());
        for study in studies {
            assert!(listed.contains(&study.to_owned()));
        }
    }

    #[tokio::test]
    async fn test_study_not_found() {
        let memory_store = Arc::new(object_store::memory::InMemory::new());
        let storage = Arc::new(
            storage::object_store_backend::test_helpers::new_for_testing(
                memory_store,
                "foxhunt-ml".to_owned(),
            ),
        );

        let config = OptunaPersistenceConfig::default();
        let persistence = OptunaPersistence::new(storage, config)
            .await
            .expect("Failed to create persistence");

        // Try to load non-existent study
        let result = persistence
            .load_study_from_minio("tuning_mamba2_nonexistent")
            .await;

        assert!(matches!(
            result,
            Err(PersistenceError::StudyNotFound { .. })
        ));
    }

    #[tokio::test]
    async fn test_delete_study() {
        let memory_store = Arc::new(object_store::memory::InMemory::new());
        let storage = Arc::new(
            storage::object_store_backend::test_helpers::new_for_testing(
                memory_store,
                "foxhunt-ml".to_owned(),
            ),
        );

        let config = OptunaPersistenceConfig::default();
        let persistence = OptunaPersistence::new(storage, config)
            .await
            .expect("Failed to create persistence");

        // Save study
        let temp_db = create_test_db().await;
        let study_name = "tuning_mamba2_delete_test";
        persistence
            .save_study_to_minio(study_name, temp_db.path())
            .await
            .expect("Failed to save study");

        // Verify it exists
        let loaded = persistence.load_study_from_minio(study_name).await;
        assert!(loaded.is_ok());

        // Delete study
        let deleted = persistence
            .delete_study(study_name)
            .await
            .expect("Failed to delete study");
        assert!(deleted);

        // Verify it's gone
        let loaded_after = persistence.load_study_from_minio(study_name).await;
        assert!(matches!(
            loaded_after,
            Err(PersistenceError::StudyNotFound { .. })
        ));
    }
}
