foxhunt/ml/src/model_registry.rs

//! Model Registry and Versioning System
//!
//! This module provides comprehensive model versioning with metadata tracking,
//! PostgreSQL storage, and API endpoints for model management.
//!
//! # Features
//!
//! - **Model Versioning**: Semantic versioning (v1.0.0) for all ML models
//! - **Metadata Tracking**: Training metrics, hyperparameters, data sources
//! - **PostgreSQL Storage**: Persistent version history with TimescaleDB
//! - **Production Tags**: Mark models as production/experimental/archived
//! - **S3 Integration**: Store model artifacts with checksums
//! - **Query API**: Query models by version, type, tag, date range
//!
//! # Example
//!
//! ```rust,no_run
//! use ml::model_registry::{ModelRegistry, ModelVersionMetadata};
//! use ml::ModelType;
//!
//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
//! let registry = ModelRegistry::new("postgresql://...", "s3://bucket/").await?;
//!
//! // Register a new model version
//! let metadata = ModelVersionMetadata {
//!     model_id: "dqn-v1.0.0".to_string(),
//!     model_type: ModelType::DQN,
//!     version: "1.0.0".to_string(),
//!     hyperparameters: serde_json::json!({
//!         "epochs": 500,
//!         "batch_size": 128,
//!         "learning_rate": 0.0001
//!     }),
//!     metrics: serde_json::json!({
//!         "final_loss": 0.001,
//!         "best_epoch": 487,
//!         "training_time_seconds": 168
//!     }),
//!     data_source: "databento_2024_Q4".to_string(),
//!     s3_location: "s3://foxhunt-ml-models/dqn/1.0.0/".to_string(),
//!     checksum: "sha256:abc123...".to_string(),
//!     training_date: chrono::Utc::now(),
//!     is_production: true,
//!     is_experimental: false,
//!     is_archived: false,
//! };
//!
//! registry.register_version(&metadata).await?;
//!
//! // Query models by version
//! let model = registry.get_model_by_version("dqn-v1.0.0").await?;
//!
//! // Get production models
//! let production_models = registry.get_production_models().await?;
//!
//! # Ok(())
//! # }
//! ```

use crate::{MLError, MLResult, ModelType};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sqlx::{postgres::PgPoolOptions, PgPool, Row};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// Checkpoint loading utilities
pub mod checkpoint_loader;

/// Model version metadata for tracking ML model versions
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelVersionMetadata {
    /// Unique model identifier (e.g., "dqn-v1.0.0")
    pub model_id: String,

    /// Type of the model
    pub model_type: ModelType,

    /// Semantic version (e.g., "1.0.0")
    pub version: String,

    /// Training date and time
    pub training_date: DateTime<Utc>,

    /// Hyperparameters used for training
    pub hyperparameters: serde_json::Value,

    /// Training and validation metrics
    pub metrics: serde_json::Value,

    /// Data source identifier (e.g., "databento_2024_Q4")
    pub data_source: String,

    /// S3 location for model artifacts
    pub s3_location: String,

    /// SHA-256 checksum of model artifacts
    pub checksum: String,

    /// Production status flag
    pub is_production: bool,

    /// Experimental status flag
    pub is_experimental: bool,

    /// Archived status flag
    pub is_archived: bool,

    /// Additional metadata
    #[serde(default)]
    pub metadata: HashMap<String, String>,

    /// Created timestamp
    #[serde(default = "Utc::now")]
    pub created_at: DateTime<Utc>,

    /// Last updated timestamp
    #[serde(default = "Utc::now")]
    pub updated_at: DateTime<Utc>,
}

impl ModelVersionMetadata {
    /// Create new model version metadata
    pub fn new(
        model_id: String,
        model_type: ModelType,
        version: String,
        data_source: String,
        s3_location: String,
    ) -> Self {
        Self {
            model_id,
            model_type,
            version,
            training_date: Utc::now(),
            hyperparameters: serde_json::json!({}),
            metrics: serde_json::json!({}),
            data_source,
            s3_location,
            checksum: String::new(),
            is_production: false,
            is_experimental: true,
            is_archived: false,
            metadata: HashMap::new(),
            created_at: Utc::now(),
            updated_at: Utc::now(),
        }
    }

    /// Mark as production model
    pub fn mark_production(mut self) -> Self {
        self.is_production = true;
        self.is_experimental = false;
        self.updated_at = Utc::now();
        self
    }

    /// Mark as experimental model
    pub fn mark_experimental(mut self) -> Self {
        self.is_production = false;
        self.is_experimental = true;
        self.updated_at = Utc::now();
        self
    }

    /// Mark as archived model
    pub fn mark_archived(mut self) -> Self {
        self.is_archived = true;
        self.updated_at = Utc::now();
        self
    }

    /// Add hyperparameter
    pub fn add_hyperparameter(&mut self, key: &str, value: serde_json::Value) {
        if let Some(obj) = self.hyperparameters.as_object_mut() {
            obj.insert(key.to_string(), value);
        }
    }

    /// Add metric
    pub fn add_metric(&mut self, key: &str, value: serde_json::Value) {
        if let Some(obj) = self.metrics.as_object_mut() {
            obj.insert(key.to_string(), value);
        }
    }

    /// Add metadata
    pub fn add_metadata(&mut self, key: &str, value: String) {
        self.metadata.insert(key.to_string(), value);
    }

    /// Set checksum
    pub fn set_checksum(&mut self, checksum: String) {
        self.checksum = checksum;
    }
}

/// Model registry for managing ML model versions
#[derive(Debug, Clone)]
pub struct ModelRegistry {
    /// PostgreSQL connection pool
    db_pool: PgPool,

    /// S3 bucket base path
    s3_base_path: String,

    /// In-memory cache for fast lookups
    cache: Arc<RwLock<HashMap<String, ModelVersionMetadata>>>,
}

impl ModelRegistry {
    /// Create new model registry
    ///
    /// # Arguments
    ///
    /// * `database_url` - PostgreSQL connection URL
    /// * `s3_base_path` - S3 bucket base path (e.g., "s3://foxhunt-ml-models/")
    ///
    /// # Returns
    ///
    /// Returns a `Result<ModelRegistry>` with the initialized registry
    ///
    /// # Errors
    ///
    /// Returns an error if database connection or schema creation fails
    pub async fn new(database_url: &str, s3_base_path: &str) -> MLResult<Self> {
        let db_pool = PgPoolOptions::new()
            .max_connections(5)
            .connect(database_url)
            .await
            .map_err(|e| MLError::ModelError(format!("Failed to connect to database: {}", e)))?;

        // Create schema if not exists
        Self::ensure_schema(&db_pool).await?;

        Ok(Self {
            db_pool,
            s3_base_path: s3_base_path.to_string(),
            cache: Arc::new(RwLock::new(HashMap::new())),
        })
    }

    /// Ensure database schema exists
    async fn ensure_schema(pool: &PgPool) -> MLResult<()> {
        // Create table
        sqlx::query(
            r#"
            CREATE TABLE IF NOT EXISTS ml_model_versions (
                id SERIAL PRIMARY KEY,
                model_id VARCHAR(255) NOT NULL UNIQUE,
                model_type VARCHAR(50) NOT NULL,
                version VARCHAR(50) NOT NULL,
                training_date TIMESTAMPTZ NOT NULL,
                hyperparameters JSONB NOT NULL DEFAULT '{}'::jsonb,
                metrics JSONB NOT NULL DEFAULT '{}'::jsonb,
                data_source VARCHAR(255) NOT NULL,
                s3_location TEXT NOT NULL,
                checksum VARCHAR(255) NOT NULL,
                is_production BOOLEAN NOT NULL DEFAULT false,
                is_experimental BOOLEAN NOT NULL DEFAULT true,
                is_archived BOOLEAN NOT NULL DEFAULT false,
                metadata JSONB NOT NULL DEFAULT '{}'::jsonb,
                created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                CONSTRAINT unique_model_version UNIQUE (model_type, version)
            )
        "#,
        )
        .execute(pool)
        .await
        .map_err(|e| MLError::ModelError(format!("Failed to create schema: {}", e)))?;

        // Create indexes (each as separate statement)
        let indexes = vec![
            "CREATE INDEX IF NOT EXISTS idx_ml_model_versions_model_type ON ml_model_versions(model_type)",
            "CREATE INDEX IF NOT EXISTS idx_ml_model_versions_version ON ml_model_versions(version)",
            "CREATE INDEX IF NOT EXISTS idx_ml_model_versions_training_date ON ml_model_versions(training_date DESC)",
            "CREATE INDEX IF NOT EXISTS idx_ml_model_versions_is_production ON ml_model_versions(is_production) WHERE is_production = true",
            "CREATE INDEX IF NOT EXISTS idx_ml_model_versions_is_experimental ON ml_model_versions(is_experimental) WHERE is_experimental = true",
            "CREATE INDEX IF NOT EXISTS idx_ml_model_versions_is_archived ON ml_model_versions(is_archived) WHERE is_archived = false",
            "CREATE INDEX IF NOT EXISTS idx_ml_model_versions_metadata_gin ON ml_model_versions USING GIN (metadata)",
            "CREATE INDEX IF NOT EXISTS idx_ml_model_versions_hyperparameters_gin ON ml_model_versions USING GIN (hyperparameters)",
            "CREATE INDEX IF NOT EXISTS idx_ml_model_versions_metrics_gin ON ml_model_versions USING GIN (metrics)",
        ];

        for index_query in indexes {
            sqlx::query(index_query)
                .execute(pool)
                .await
                .map_err(|e| MLError::ModelError(format!("Failed to create index: {}", e)))?;
        }

        Ok(())
    }

    /// Register a new model version
    pub async fn register_version(&self, metadata: &ModelVersionMetadata) -> MLResult<()> {
        let query = r#"
            INSERT INTO ml_model_versions (
                model_id, model_type, version, training_date,
                hyperparameters, metrics, data_source, s3_location,
                checksum, is_production, is_experimental, is_archived, metadata
            ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)
            ON CONFLICT (model_id)
            DO UPDATE SET
                training_date = EXCLUDED.training_date,
                hyperparameters = EXCLUDED.hyperparameters,
                metrics = EXCLUDED.metrics,
                data_source = EXCLUDED.data_source,
                s3_location = EXCLUDED.s3_location,
                checksum = EXCLUDED.checksum,
                is_production = EXCLUDED.is_production,
                is_experimental = EXCLUDED.is_experimental,
                is_archived = EXCLUDED.is_archived,
                metadata = EXCLUDED.metadata,
                updated_at = NOW()
        "#;

        let model_type_str = format!("{:?}", metadata.model_type);

        sqlx::query(query)
            .bind(&metadata.model_id)
            .bind(&model_type_str)
            .bind(&metadata.version)
            .bind(metadata.training_date)
            .bind(&metadata.hyperparameters)
            .bind(&metadata.metrics)
            .bind(&metadata.data_source)
            .bind(&metadata.s3_location)
            .bind(&metadata.checksum)
            .bind(metadata.is_production)
            .bind(metadata.is_experimental)
            .bind(metadata.is_archived)
            .bind(serde_json::to_value(&metadata.metadata).unwrap_or_default())
            .execute(&self.db_pool)
            .await
            .map_err(|e| MLError::ModelError(format!("Failed to register version: {}", e)))?;

        // Update cache
        self.cache
            .write()
            .await
            .insert(metadata.model_id.clone(), metadata.clone());

        tracing::info!("Registered model version: {}", metadata.model_id);

        Ok(())
    }

    /// Get model by version
    pub async fn get_model_by_version(&self, model_id: &str) -> MLResult<ModelVersionMetadata> {
        // Check cache first
        if let Some(metadata) = self.cache.read().await.get(model_id) {
            return Ok(metadata.clone());
        }

        // Query database
        let query = r#"
            SELECT model_id, model_type, version, training_date,
                   hyperparameters, metrics, data_source, s3_location,
                   checksum, is_production, is_experimental, is_archived,
                   metadata, created_at, updated_at
            FROM ml_model_versions
            WHERE model_id = $1
        "#;

        let row = sqlx::query(query)
            .bind(model_id)
            .fetch_one(&self.db_pool)
            .await
            .map_err(|e| MLError::ModelNotFound(format!("Model {} not found: {}", model_id, e)))?;

        let metadata = self.row_to_metadata(row)?;

        // Update cache
        self.cache
            .write()
            .await
            .insert(model_id.to_string(), metadata.clone());

        Ok(metadata)
    }

    /// Get all production models
    pub async fn get_production_models(&self) -> MLResult<Vec<ModelVersionMetadata>> {
        let query = r#"
            SELECT model_id, model_type, version, training_date,
                   hyperparameters, metrics, data_source, s3_location,
                   checksum, is_production, is_experimental, is_archived,
                   metadata, created_at, updated_at
            FROM ml_model_versions
            WHERE is_production = true AND is_archived = false
            ORDER BY training_date DESC
        "#;

        let rows = sqlx::query(query)
            .fetch_all(&self.db_pool)
            .await
            .map_err(|e| {
                MLError::ModelError(format!("Failed to query production models: {}", e))
            })?;

        let mut models = Vec::new();
        for row in rows {
            models.push(self.row_to_metadata(row)?);
        }

        Ok(models)
    }

    /// Get all experimental models
    pub async fn get_experimental_models(&self) -> MLResult<Vec<ModelVersionMetadata>> {
        let query = r#"
            SELECT model_id, model_type, version, training_date,
                   hyperparameters, metrics, data_source, s3_location,
                   checksum, is_production, is_experimental, is_archived,
                   metadata, created_at, updated_at
            FROM ml_model_versions
            WHERE is_experimental = true AND is_archived = false
            ORDER BY training_date DESC
        "#;

        let rows = sqlx::query(query)
            .fetch_all(&self.db_pool)
            .await
            .map_err(|e| {
                MLError::ModelError(format!("Failed to query experimental models: {}", e))
            })?;

        let mut models = Vec::new();
        for row in rows {
            models.push(self.row_to_metadata(row)?);
        }

        Ok(models)
    }

    /// Get models by type
    pub async fn get_models_by_type(
        &self,
        model_type: ModelType,
    ) -> MLResult<Vec<ModelVersionMetadata>> {
        let query = r#"
            SELECT model_id, model_type, version, training_date,
                   hyperparameters, metrics, data_source, s3_location,
                   checksum, is_production, is_experimental, is_archived,
                   metadata, created_at, updated_at
            FROM ml_model_versions
            WHERE model_type = $1 AND is_archived = false
            ORDER BY training_date DESC
        "#;

        let model_type_str = format!("{:?}", model_type);

        let rows = sqlx::query(query)
            .bind(&model_type_str)
            .fetch_all(&self.db_pool)
            .await
            .map_err(|e| MLError::ModelError(format!("Failed to query models by type: {}", e)))?;

        let mut models = Vec::new();
        for row in rows {
            models.push(self.row_to_metadata(row)?);
        }

        Ok(models)
    }

    /// Get models by date range
    pub async fn get_models_by_date_range(
        &self,
        start_date: DateTime<Utc>,
        end_date: DateTime<Utc>,
    ) -> MLResult<Vec<ModelVersionMetadata>> {
        let query = r#"
            SELECT model_id, model_type, version, training_date,
                   hyperparameters, metrics, data_source, s3_location,
                   checksum, is_production, is_experimental, is_archived,
                   metadata, created_at, updated_at
            FROM ml_model_versions
            WHERE training_date >= $1 AND training_date <= $2
            ORDER BY training_date DESC
        "#;

        let rows = sqlx::query(query)
            .bind(start_date)
            .bind(end_date)
            .fetch_all(&self.db_pool)
            .await
            .map_err(|e| MLError::ModelError(format!("Failed to query models by date: {}", e)))?;

        let mut models = Vec::new();
        for row in rows {
            models.push(self.row_to_metadata(row)?);
        }

        Ok(models)
    }

    /// Mark model as production
    pub async fn mark_production(&self, model_id: &str) -> MLResult<()> {
        let query = r#"
            UPDATE ml_model_versions
            SET is_production = true, is_experimental = false, updated_at = NOW()
            WHERE model_id = $1
        "#;

        sqlx::query(query)
            .bind(model_id)
            .execute(&self.db_pool)
            .await
            .map_err(|e| MLError::ModelError(format!("Failed to mark production: {}", e)))?;

        // Invalidate cache
        self.cache.write().await.remove(model_id);

        tracing::info!("Marked model as production: {}", model_id);

        Ok(())
    }

    /// Mark model as experimental
    pub async fn mark_experimental(&self, model_id: &str) -> MLResult<()> {
        let query = r#"
            UPDATE ml_model_versions
            SET is_production = false, is_experimental = true, updated_at = NOW()
            WHERE model_id = $1
        "#;

        sqlx::query(query)
            .bind(model_id)
            .execute(&self.db_pool)
            .await
            .map_err(|e| MLError::ModelError(format!("Failed to mark experimental: {}", e)))?;

        // Invalidate cache
        self.cache.write().await.remove(model_id);

        tracing::info!("Marked model as experimental: {}", model_id);

        Ok(())
    }

    /// Archive model
    pub async fn archive_model(&self, model_id: &str) -> MLResult<()> {
        let query = r#"
            UPDATE ml_model_versions
            SET is_archived = true, updated_at = NOW()
            WHERE model_id = $1
        "#;

        sqlx::query(query)
            .bind(model_id)
            .execute(&self.db_pool)
            .await
            .map_err(|e| MLError::ModelError(format!("Failed to archive model: {}", e)))?;

        // Invalidate cache
        self.cache.write().await.remove(model_id);

        tracing::info!("Archived model: {}", model_id);

        Ok(())
    }

    /// Delete model version (permanent)
    pub async fn delete_version(&self, model_id: &str) -> MLResult<()> {
        let query = r#"
            DELETE FROM ml_model_versions
            WHERE model_id = $1
        "#;

        sqlx::query(query)
            .bind(model_id)
            .execute(&self.db_pool)
            .await
            .map_err(|e| MLError::ModelError(format!("Failed to delete version: {}", e)))?;

        // Invalidate cache
        self.cache.write().await.remove(model_id);

        tracing::warn!("Deleted model version: {}", model_id);

        Ok(())
    }

    /// Get registry statistics
    pub async fn get_statistics(&self) -> MLResult<RegistryStatistics> {
        let query = r#"
            SELECT
                COUNT(*) FILTER (WHERE is_production = true AND is_archived = false) as production_count,
                COUNT(*) FILTER (WHERE is_experimental = true AND is_archived = false) as experimental_count,
                COUNT(*) FILTER (WHERE is_archived = true) as archived_count,
                COUNT(*) as total_count,
                COUNT(DISTINCT model_type) as model_types_count,
                MAX(training_date) as latest_training_date,
                MIN(training_date) as earliest_training_date
            FROM ml_model_versions
        "#;

        let row = sqlx::query(query)
            .fetch_one(&self.db_pool)
            .await
            .map_err(|e| MLError::ModelError(format!("Failed to get statistics: {}", e)))?;

        Ok(RegistryStatistics {
            production_count: row.try_get::<i64, _>("production_count").unwrap_or(0) as usize,
            experimental_count: row.try_get::<i64, _>("experimental_count").unwrap_or(0) as usize,
            archived_count: row.try_get::<i64, _>("archived_count").unwrap_or(0) as usize,
            total_count: row.try_get::<i64, _>("total_count").unwrap_or(0) as usize,
            model_types_count: row.try_get::<i64, _>("model_types_count").unwrap_or(0) as usize,
            latest_training_date: row.try_get("latest_training_date").ok(),
            earliest_training_date: row.try_get("earliest_training_date").ok(),
        })
    }

    /// Convert database row to metadata
    fn row_to_metadata(&self, row: sqlx::postgres::PgRow) -> MLResult<ModelVersionMetadata> {
        let model_type_str: String = row
            .try_get("model_type")
            .map_err(|e| MLError::ModelError(format!("Failed to get model_type: {}", e)))?;

        let model_type = ModelType::from_str(&model_type_str).ok_or_else(|| {
            MLError::ModelError(format!("Invalid model type: {}", model_type_str))
        })?;

        let metadata_json: serde_json::Value = row
            .try_get("metadata")
            .map_err(|e| MLError::ModelError(format!("Failed to get metadata: {}", e)))?;

        let metadata_map: HashMap<String, String> =
            serde_json::from_value(metadata_json).unwrap_or_default();

        Ok(ModelVersionMetadata {
            model_id: row
                .try_get("model_id")
                .map_err(|e| MLError::ModelError(format!("Failed to get model_id: {}", e)))?,
            model_type,
            version: row
                .try_get("version")
                .map_err(|e| MLError::ModelError(format!("Failed to get version: {}", e)))?,
            training_date: row
                .try_get("training_date")
                .map_err(|e| MLError::ModelError(format!("Failed to get training_date: {}", e)))?,
            hyperparameters: row.try_get("hyperparameters").map_err(|e| {
                MLError::ModelError(format!("Failed to get hyperparameters: {}", e))
            })?,
            metrics: row
                .try_get("metrics")
                .map_err(|e| MLError::ModelError(format!("Failed to get metrics: {}", e)))?,
            data_source: row
                .try_get("data_source")
                .map_err(|e| MLError::ModelError(format!("Failed to get data_source: {}", e)))?,
            s3_location: row
                .try_get("s3_location")
                .map_err(|e| MLError::ModelError(format!("Failed to get s3_location: {}", e)))?,
            checksum: row
                .try_get("checksum")
                .map_err(|e| MLError::ModelError(format!("Failed to get checksum: {}", e)))?,
            is_production: row
                .try_get("is_production")
                .map_err(|e| MLError::ModelError(format!("Failed to get is_production: {}", e)))?,
            is_experimental: row.try_get("is_experimental").map_err(|e| {
                MLError::ModelError(format!("Failed to get is_experimental: {}", e))
            })?,
            is_archived: row
                .try_get("is_archived")
                .map_err(|e| MLError::ModelError(format!("Failed to get is_archived: {}", e)))?,
            metadata: metadata_map,
            created_at: row
                .try_get("created_at")
                .map_err(|e| MLError::ModelError(format!("Failed to get created_at: {}", e)))?,
            updated_at: row
                .try_get("updated_at")
                .map_err(|e| MLError::ModelError(format!("Failed to get updated_at: {}", e)))?,
        })
    }
}

/// Registry statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RegistryStatistics {
    /// Number of production models
    pub production_count: usize,
    /// Number of experimental models
    pub experimental_count: usize,
    /// Number of archived models
    pub archived_count: usize,
    /// Total number of models
    pub total_count: usize,
    /// Number of distinct model types
    pub model_types_count: usize,
    /// Latest training date
    pub latest_training_date: Option<DateTime<Utc>>,
    /// Earliest training date
    pub earliest_training_date: Option<DateTime<Utc>>,
}

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

    #[tokio::test]
    #[ignore = "Requires PostgreSQL"]
    async fn test_model_registry_new() {
        let registry = ModelRegistry::new(
            "postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt",
            "s3://foxhunt-ml-models/",
        )
        .await;

        assert!(registry.is_ok());
    }

    #[tokio::test]
    #[ignore = "Requires PostgreSQL"]
    async fn test_register_and_retrieve_model() {
        let registry = ModelRegistry::new(
            "postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt",
            "s3://foxhunt-ml-models/",
        )
        .await
        .unwrap();

        let mut metadata = ModelVersionMetadata::new(
            "dqn-test-v1.0.0".to_string(),
            ModelType::DQN,
            "1.0.0".to_string(),
            "test_data".to_string(),
            "s3://foxhunt-ml-models/dqn/1.0.0/".to_string(),
        );

        metadata.add_hyperparameter("epochs", serde_json::json!(500));
        metadata.add_metric("final_loss", serde_json::json!(0.001));
        metadata.set_checksum("sha256:test123".to_string());

        // Register
        registry.register_version(&metadata).await.unwrap();

        // Retrieve
        let retrieved = registry
            .get_model_by_version("dqn-test-v1.0.0")
            .await
            .unwrap();

        assert_eq!(retrieved.model_id, "dqn-test-v1.0.0");
        assert_eq!(retrieved.version, "1.0.0");
    }
}