foxhunt/ml-data/src/models.rs

//! Model Artifacts Repository
//! 
//! Manages ML model artifacts, versioning, metadata, and deployment lifecycle
//! for HFT trading systems with PostgreSQL integration.

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::{MlDataError, Result};
use database::{DatabasePool, DatabaseConnection};

/// Model artifacts repository for ML model lifecycle management
#[derive(Clone)]
pub struct ModelRepository {
    pool: DatabasePool,
    storage_path: PathBuf,
}

impl ModelRepository {
    pub async fn new(pool: DatabasePool, storage_path: String) -> Result<Self> {
        let storage_path = PathBuf::from(storage_path);
        
        // Ensure storage directory exists
        if !storage_path.exists() {
            std::fs::create_dir_all(&storage_path)?;
        }
        
        let repo = Self { pool, storage_path };
        repo.initialize_schema().await?;
        Ok(repo)
    }
    
    /// Initialize database schema for model artifacts
    pub async fn initialize_schema(&self) -> Result<()> {
        let conn = self.pool.get().await?;
        
        // Model versions table
        conn.execute(r#"
            CREATE TABLE IF NOT EXISTS ml_model_versions (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                model_name VARCHAR NOT NULL,
                version VARCHAR NOT NULL,
                model_type VARCHAR NOT NULL,
                framework VARCHAR NOT NULL,
                created_at TIMESTAMPTZ DEFAULT NOW(),
                updated_at TIMESTAMPTZ DEFAULT NOW(),
                created_by VARCHAR NOT NULL,
                status model_status DEFAULT 'training',
                deployment_status deployment_status DEFAULT 'not_deployed',
                file_path VARCHAR NOT NULL,
                file_size BIGINT NOT NULL,
                checksum VARCHAR NOT NULL,
                metadata JSONB DEFAULT '{}',
                training_config JSONB DEFAULT '{}',
                performance_metrics JSONB DEFAULT '{}',
                UNIQUE(model_name, version)
            )
        "#).await?;
        
        // Create enums
        conn.execute(r#"
            DO $$ BEGIN
                CREATE TYPE model_status AS ENUM ('training', 'trained', 'validated', 'deployed', 'deprecated');
            EXCEPTION
                WHEN duplicate_object THEN null;
            END $$;
        "#).await?;
        
        conn.execute(r#"
            DO $$ BEGIN
                CREATE TYPE deployment_status AS ENUM ('not_deployed', 'staging', 'production', 'canary', 'rollback');
            EXCEPTION
                WHEN duplicate_object THEN null;
            END $$;
        "#).await?;
        
        // Model dependencies table (for ensemble models)
        conn.execute(r#"
            CREATE TABLE IF NOT EXISTS ml_model_dependencies (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                parent_model_id UUID NOT NULL REFERENCES ml_model_versions(id) ON DELETE CASCADE,
                dependency_model_id UUID NOT NULL REFERENCES ml_model_versions(id) ON DELETE CASCADE,
                dependency_type VARCHAR NOT NULL,
                weight DOUBLE PRECISION DEFAULT 1.0,
                created_at TIMESTAMPTZ DEFAULT NOW()
            )
        "#).await?;
        
        // Model deployment history
        conn.execute(r#"
            CREATE TABLE IF NOT EXISTS ml_model_deployments (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                model_id UUID NOT NULL REFERENCES ml_model_versions(id) ON DELETE CASCADE,
                environment VARCHAR NOT NULL,
                deployment_status deployment_status NOT NULL,
                deployed_at TIMESTAMPTZ DEFAULT NOW(),
                deployed_by VARCHAR NOT NULL,
                rollback_model_id UUID REFERENCES ml_model_versions(id),
                deployment_config JSONB DEFAULT '{}',
                health_check_url VARCHAR,
                notes TEXT
            )
        "#).await?;
        
        // Indexes
        conn.execute("CREATE INDEX IF NOT EXISTS idx_models_name_version ON ml_model_versions(model_name, version)").await?;
        conn.execute("CREATE INDEX IF NOT EXISTS idx_models_status ON ml_model_versions(status)").await?;
        conn.execute("CREATE INDEX IF NOT EXISTS idx_deployments_environment ON ml_model_deployments(environment, deployment_status)").await?;
        
        Ok(())
    }
    
    /// Save a new model artifact
    pub async fn save_model(&self, request: SaveModelRequest) -> Result<ModelArtifact> {
        let mut conn = self.pool.get().await?;
        let tx = conn.begin().await?;
        
        // Validate model request
        self.validate_save_request(&request).await?;
        
        // Check for version conflicts
        if self.model_version_exists(&request.model_name, &request.version).await? {
            return Err(MlDataError::VersionConflict {
                message: format!("Model {} version {} already exists", 
                    request.model_name, request.version)
            });
        }
        
        let model_id = Uuid::new_v4();
        
        // Save model file to storage
        let file_path = self.get_model_file_path(&request.model_name, &request.version);
        std::fs::write(&file_path, &request.model_data)?;
        
        // Calculate file checksum
        let checksum = self.calculate_checksum(&request.model_data);
        let file_size = request.model_data.len() as i64;
        
        // Insert model record
        tx.execute(
            r#"INSERT INTO ml_model_versions 
               (id, model_name, version, model_type, framework, created_by, 
                file_path, file_size, checksum, metadata, training_config)
               VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)"#,
            &[&model_id, &request.model_name, &request.version, 
              &request.model_type, &request.framework, &request.created_by,
              &file_path.to_string_lossy().to_string(), &file_size, &checksum,
              &request.metadata, &request.training_config]
        ).await?;
        
        tx.commit().await?;
        
        let artifact = ModelArtifact {
            id: model_id,
            model_name: request.model_name,
            version: request.version,
            model_type: request.model_type,
            framework: request.framework,
            created_at: Utc::now(),
            updated_at: Utc::now(),
            created_by: request.created_by,
            status: ModelStatus::Trained,
            deployment_status: DeploymentStatus::NotDeployed,
            file_path: file_path.to_string_lossy().to_string(),
            file_size,
            checksum,
            metadata: request.metadata,
            training_config: request.training_config,
            performance_metrics: serde_json::Value::Object(serde_json::Map::new()),
            dependencies: Vec::new(),
        };
        
        tracing::info!("Saved model artifact: {} v{}", artifact.model_name, artifact.version);
        Ok(artifact)
    }
    
    /// Load a model artifact by name and version
    pub async fn load_model(&self, model_name: &str, version: &str) -> Result<ModelArtifact> {
        let conn = self.pool.get().await?;
        
        let row = conn.query_one(
            r#"SELECT id, model_name, version, model_type, framework, created_at, updated_at,
                      created_by, status, deployment_status, file_path, file_size, checksum,
                      metadata, training_config, performance_metrics
               FROM ml_model_versions 
               WHERE model_name = $1 AND version = $2"#,
            &[&model_name, &version]
        ).await.map_err(|_| MlDataError::NotFound {
            resource_type: "Model".to_string(),
            id: format!("{}:{}", model_name, version),
        })?;
        
        let model_id: Uuid = row.get("id");
        let dependencies = self.load_model_dependencies(model_id).await?;
        
        Ok(ModelArtifact {
            id: model_id,
            model_name: row.get("model_name"),
            version: row.get("version"),
            model_type: row.get("model_type"),
            framework: row.get("framework"),
            created_at: row.get("created_at"),
            updated_at: row.get("updated_at"),
            created_by: row.get("created_by"),
            status: self.parse_model_status(row.get("status"))?,
            deployment_status: self.parse_deployment_status(row.get("deployment_status"))?,
            file_path: row.get("file_path"),
            file_size: row.get("file_size"),
            checksum: row.get("checksum"),
            metadata: row.get("metadata"),
            training_config: row.get("training_config"),
            performance_metrics: row.get("performance_metrics"),
            dependencies,
        })
    }
    
    /// Load model binary data
    pub async fn load_model_data(&self, model_name: &str, version: &str) -> Result<Vec<u8>> {
        let artifact = self.load_model(model_name, version).await?;
        let data = std::fs::read(&artifact.file_path)?;
        
        // Verify checksum
        let calculated_checksum = self.calculate_checksum(&data);
        if calculated_checksum != artifact.checksum {
            return Err(MlDataError::Validation {
                message: "Model file checksum mismatch - data may be corrupted".to_string()
            });
        }
        
        Ok(data)
    }
    
    /// Update model status
    pub async fn update_status(&self, model_id: Uuid, status: ModelStatus) -> Result<()> {
        let conn = self.pool.get().await?;
        
        conn.execute(
            "UPDATE ml_model_versions SET status = $1, updated_at = NOW() WHERE id = $2",
            &[&status.to_string(), &model_id]
        ).await?;
        
        tracing::info!("Updated model {} status to {:?}", model_id, status);
        Ok(())
    }
    
    /// Deploy a model to an environment
    pub async fn deploy_model(&self, request: DeployModelRequest) -> Result<DeploymentRecord> {
        let mut conn = self.pool.get().await?;
        let tx = conn.begin().await?;
        
        let deployment_id = Uuid::new_v4();
        
        // Update model deployment status
        tx.execute(
            "UPDATE ml_model_versions SET deployment_status = $1, updated_at = NOW() WHERE id = $2",
            &[&request.deployment_status.to_string(), &request.model_id]
        ).await?;
        
        // Record deployment
        tx.execute(
            r#"INSERT INTO ml_model_deployments 
               (id, model_id, environment, deployment_status, deployed_by, 
                rollback_model_id, deployment_config, health_check_url, notes)
               VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)"#,
            &[&deployment_id, &request.model_id, &request.environment,
              &request.deployment_status.to_string(), &request.deployed_by,
              &request.rollback_model_id, &request.deployment_config,
              &request.health_check_url, &request.notes]
        ).await?;
        
        tx.commit().await?;
        
        let record = DeploymentRecord {
            id: deployment_id,
            model_id: request.model_id,
            environment: request.environment,
            deployment_status: request.deployment_status,
            deployed_at: Utc::now(),
            deployed_by: request.deployed_by,
            rollback_model_id: request.rollback_model_id,
            deployment_config: request.deployment_config,
            health_check_url: request.health_check_url,
            notes: request.notes,
        };
        
        tracing::info!("Deployed model {} to {}", request.model_id, record.environment);
        Ok(record)
    }
    
    /// List all versions of a model
    pub async fn list_model_versions(&self, model_name: &str) -> Result<Vec<ModelVersion>> {
        let conn = self.pool.get().await?;
        
        let rows = conn.query(
            r#"SELECT version, status, deployment_status, created_at, file_size,
                      performance_metrics
               FROM ml_model_versions 
               WHERE model_name = $1 
               ORDER BY created_at DESC"#,
            &[&model_name]
        ).await?;
        
        let mut versions = Vec::new();
        for row in rows {
            versions.push(ModelVersion {
                version: row.get("version"),
                status: self.parse_model_status(row.get("status"))?,
                deployment_status: self.parse_deployment_status(row.get("deployment_status"))?,
                created_at: row.get("created_at"),
                file_size: row.get("file_size"),
                performance_metrics: row.get("performance_metrics"),
            });
        }
        
        Ok(versions)
    }
    
    /// Add model dependencies (for ensemble models)
    pub async fn add_dependencies(
        &self, 
        parent_model_id: Uuid, 
        dependencies: Vec<ModelDependency>
    ) -> Result<()> {
        let mut conn = self.pool.get().await?;
        let tx = conn.begin().await?;
        
        for dep in dependencies {
            tx.execute(
                r#"INSERT INTO ml_model_dependencies 
                   (parent_model_id, dependency_model_id, dependency_type, weight)
                   VALUES ($1, $2, $3, $4)"#,
                &[&parent_model_id, &dep.model_id, &dep.dependency_type, &dep.weight]
            ).await?;
        }
        
        tx.commit().await?;
        tracing::info!("Added dependencies for model {}", parent_model_id);
        Ok(())
    }
    
    /// Load model dependencies
    async fn load_model_dependencies(&self, model_id: Uuid) -> Result<Vec<ModelDependency>> {
        let conn = self.pool.get().await?;
        
        let rows = conn.query(
            r#"SELECT dependency_model_id, dependency_type, weight
               FROM ml_model_dependencies 
               WHERE parent_model_id = $1"#,
            &[&model_id]
        ).await?;
        
        let mut dependencies = Vec::new();
        for row in rows {
            dependencies.push(ModelDependency {
                model_id: row.get("dependency_model_id"),
                dependency_type: row.get("dependency_type"),
                weight: row.get("weight"),
            });
        }
        
        Ok(dependencies)
    }
    
    /// Generate file path for model artifact
    fn get_model_file_path(&self, model_name: &str, version: &str) -> PathBuf {
        self.storage_path
            .join(model_name)
            .join(format!("{}.model", version))
    }
    
    /// Calculate SHA-256 checksum
    fn calculate_checksum(&self, data: &[u8]) -> String {
        use sha2::{Sha256, Digest};
        let mut hasher = Sha256::new();
        hasher.update(data);
        format!("{:x}", hasher.finalize())
    }
    
    /// Validate save model request
    async fn validate_save_request(&self, request: &SaveModelRequest) -> Result<()> {
        if request.model_name.trim().is_empty() {
            return Err(MlDataError::Validation {
                message: "Model name cannot be empty".to_string()
            });
        }
        
        if request.version.trim().is_empty() {
            return Err(MlDataError::Validation {
                message: "Model version cannot be empty".to_string()
            });
        }
        
        if request.model_data.is_empty() {
            return Err(MlDataError::Validation {
                message: "Model data cannot be empty".to_string()
            });
        }
        
        Ok(())
    }
    
    /// Check if model version exists
    async fn model_version_exists(&self, model_name: &str, version: &str) -> Result<bool> {
        let conn = self.pool.get().await?;
        let count: i64 = conn.query_one(
            "SELECT COUNT(*) FROM ml_model_versions WHERE model_name = $1 AND version = $2",
            &[&model_name, &version]
        ).await?.get(0);
        
        Ok(count > 0)
    }
    
    /// Parse model status from database
    fn parse_model_status(&self, status_str: &str) -> Result<ModelStatus> {
        match status_str {
            "training" => Ok(ModelStatus::Training),
            "trained" => Ok(ModelStatus::Trained),
            "validated" => Ok(ModelStatus::Validated),
            "deployed" => Ok(ModelStatus::Deployed),
            "deprecated" => Ok(ModelStatus::Deprecated),
            _ => Err(MlDataError::Validation {
                message: format!("Invalid model status: {}", status_str)
            })
        }
    }
    
    /// Parse deployment status from database
    fn parse_deployment_status(&self, status_str: &str) -> Result<DeploymentStatus> {
        match status_str {
            "not_deployed" => Ok(DeploymentStatus::NotDeployed),
            "staging" => Ok(DeploymentStatus::Staging),
            "production" => Ok(DeploymentStatus::Production),
            "canary" => Ok(DeploymentStatus::Canary),
            "rollback" => Ok(DeploymentStatus::Rollback),
            _ => Err(MlDataError::Validation {
                message: format!("Invalid deployment status: {}", status_str)
            })
        }
    }
    
    /// Health check for model repository
    pub async fn health_check(&self) -> Result<bool> {
        let conn = self.pool.get().await?;
        let _: i64 = conn.query_one("SELECT COUNT(*) FROM ml_model_versions", &[]).await?;
        Ok(self.storage_path.exists())
    }
}

/// Request to save a new model artifact
#[derive(Debug)]
pub struct SaveModelRequest {
    pub model_name: String,
    pub version: String,
    pub model_type: String,
    pub framework: String,
    pub created_by: String,
    pub model_data: Vec<u8>,
    pub metadata: serde_json::Value,
    pub training_config: serde_json::Value,
}

/// Model artifact representation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelArtifact {
    pub id: Uuid,
    pub model_name: String,
    pub version: String,
    pub model_type: String,
    pub framework: String,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
    pub created_by: String,
    pub status: ModelStatus,
    pub deployment_status: DeploymentStatus,
    pub file_path: String,
    pub file_size: i64,
    pub checksum: String,
    pub metadata: serde_json::Value,
    pub training_config: serde_json::Value,
    pub performance_metrics: serde_json::Value,
    pub dependencies: Vec<ModelDependency>,
}

/// Model status enumeration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ModelStatus {
    Training,
    Trained,
    Validated,
    Deployed,
    Deprecated,
}

impl ToString for ModelStatus {
    fn to_string(&self) -> String {
        match self {
            ModelStatus::Training => "training".to_string(),
            ModelStatus::Trained => "trained".to_string(),
            ModelStatus::Validated => "validated".to_string(),
            ModelStatus::Deployed => "deployed".to_string(),
            ModelStatus::Deprecated => "deprecated".to_string(),
        }
    }
}

/// Deployment status enumeration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum DeploymentStatus {
    NotDeployed,
    Staging,
    Production,
    Canary,
    Rollback,
}

impl ToString for DeploymentStatus {
    fn to_string(&self) -> String {
        match self {
            DeploymentStatus::NotDeployed => "not_deployed".to_string(),
            DeploymentStatus::Staging => "staging".to_string(),
            DeploymentStatus::Production => "production".to_string(),
            DeploymentStatus::Canary => "canary".to_string(),
            DeploymentStatus::Rollback => "rollback".to_string(),
        }
    }
}

/// Model version information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelVersion {
    pub version: String,
    pub status: ModelStatus,
    pub deployment_status: DeploymentStatus,
    pub created_at: DateTime<Utc>,
    pub file_size: i64,
    pub performance_metrics: serde_json::Value,
}

/// Model dependency for ensemble models
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelDependency {
    pub model_id: Uuid,
    pub dependency_type: String,
    pub weight: f64,
}

/// Request to deploy a model
#[derive(Debug)]
pub struct DeployModelRequest {
    pub model_id: Uuid,
    pub environment: String,
    pub deployment_status: DeploymentStatus,
    pub deployed_by: String,
    pub rollback_model_id: Option<Uuid>,
    pub deployment_config: serde_json::Value,
    pub health_check_url: Option<String>,
    pub notes: Option<String>,
}

/// Deployment record
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DeploymentRecord {
    pub id: Uuid,
    pub model_id: Uuid,
    pub environment: String,
    pub deployment_status: DeploymentStatus,
    pub deployed_at: DateTime<Utc>,
    pub deployed_by: String,
    pub rollback_model_id: Option<Uuid>,
    pub deployment_config: serde_json::Value,
    pub health_check_url: Option<String>,
    pub notes: Option<String>,
}

/// Model metadata for lightweight operations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelMetadata {
    pub id: Uuid,
    pub name: String,
    pub version: String,
    pub model_type: String,
    pub framework: String,
    pub status: ModelStatus,
    pub deployment_status: DeploymentStatus,
    pub created_at: DateTime<Utc>,
    pub file_size: i64,
}