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 chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use uuid::Uuid;

use crate::{MlDataError, Result};
use database::Database;

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

impl ModelRepository {
    pub async fn new(db: Database, 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 { db, storage_path };
        repo.initialize_schema().await?;
        Ok(repo)
    }

    /// Initialize database schema for model artifacts
    pub async fn initialize_schema(&self) -> Result<()> {
        // Model versions table
        self.db
            .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
        self.db.execute(r#"
            DO $$ BEGIN
                CREATE TYPE model_status AS ENUM ('training', 'trained', 'validated', 'deployed', 'deprecated');
            EXCEPTION
                WHEN duplicate_object THEN null;
            END $$;
        "#).await?;

        self.db.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)
        self.db.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
        self.db
            .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
        self.db.execute("CREATE INDEX IF NOT EXISTS idx_models_name_version ON ml_model_versions(model_name, version)").await?;
        self.db
            .execute("CREATE INDEX IF NOT EXISTS idx_models_status ON ml_model_versions(status)")
            .await?;
        self.db.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 tx = self.db.begin_transaction().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
        let query = format!(
            r#"INSERT INTO ml_model_versions
               (id, model_name, version, model_type, framework, created_by,
                file_path, file_size, checksum, metadata, training_config)
               VALUES ('{}', '{}', '{}', '{}', '{}', '{}', '{}', {}, '{}', '{}', '{}')"#,
            model_id,
            request.model_name.replace("'", "''"),
            request.version.replace("'", "''"),
            request.model_type.replace("'", "''"),
            request.framework.replace("'", "''"),
            request.created_by.replace("'", "''"),
            file_path.to_string_lossy().replace("'", "''"),
            file_size,
            checksum.replace("'", "''"),
            request.metadata.to_string().replace("'", "''"),
            request.training_config.to_string().replace("'", "''")
        );
        tx.execute(&query).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 mut conn = self.db.acquire().await?;

        let row = sqlx::query_as::<
            _,
            (
                Uuid,
                String,
                String,
                String,
                String,
                DateTime<Utc>,
                DateTime<Utc>,
                String,
                String,
                String,
                String,
                i64,
                String,
                serde_json::Value,
                serde_json::Value,
                serde_json::Value,
            ),
        >(
            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"#,
        )
        .bind(model_name)
        .bind(version)
        .fetch_one(conn.as_mut())
        .await
        .map_err(|_| MlDataError::NotFound {
            resource_type: "Model".to_string(),
            id: format!("{}:{}", model_name, version),
        })?;

        let (
            model_id,
            model_name,
            version,
            model_type,
            framework,
            created_at,
            updated_at,
            created_by,
            status_str,
            deployment_status_str,
            file_path,
            file_size,
            checksum,
            metadata,
            training_config,
            performance_metrics,
        ) = row;
        let dependencies = self.load_model_dependencies(model_id).await?;

        Ok(ModelArtifact {
            id: model_id,
            model_name,
            version,
            model_type,
            framework,
            created_at,
            updated_at,
            created_by,
            status: self.parse_model_status(&status_str)?,
            deployment_status: self.parse_deployment_status(&deployment_status_str)?,
            file_path,
            file_size,
            checksum,
            metadata,
            training_config,
            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 mut conn = self.db.acquire().await?;

        sqlx::query("UPDATE ml_model_versions SET status = $1, updated_at = NOW() WHERE id = $2")
            .bind(status.to_string())
            .bind(model_id)
            .execute(conn.as_mut())
            .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 tx = self.db.begin_transaction().await?;

        let deployment_id = Uuid::new_v4();

        // Update model deployment status
        let update_query = format!(
            "UPDATE ml_model_versions SET deployment_status = '{}', updated_at = NOW() WHERE id = '{}'",
            request.deployment_status.to_string().replace("'", "''"),
            request.model_id
        );
        tx.execute(&update_query).await?;

        // Record deployment
        let rollback_model_id = request
            .rollback_model_id
            .map(|id| format!("'{}'", id))
            .unwrap_or("NULL".to_string());
        let health_check_url = request
            .health_check_url
            .as_ref()
            .map(|url| format!("'{}'", url.replace("'", "''")))
            .unwrap_or("NULL".to_string());
        let notes = request
            .notes
            .as_ref()
            .map(|n| format!("'{}'", n.replace("'", "''")))
            .unwrap_or("NULL".to_string());
        let insert_query = format!(
            r#"INSERT INTO ml_model_deployments
               (id, model_id, environment, deployment_status, deployed_by,
                rollback_model_id, deployment_config, health_check_url, notes)
               VALUES ('{}', '{}', '{}', '{}', '{}', {}, '{}', {}, {})"#,
            deployment_id,
            request.model_id,
            request.environment.replace("'", "''"),
            request.deployment_status.to_string().replace("'", "''"),
            request.deployed_by.replace("'", "''"),
            rollback_model_id,
            request.deployment_config.to_string().replace("'", "''"),
            health_check_url,
            notes
        );
        tx.execute(&insert_query).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 mut conn = self.db.acquire().await?;

        let rows = sqlx::query_as::<
            _,
            (
                String,
                String,
                String,
                DateTime<Utc>,
                i64,
                serde_json::Value,
            ),
        >(
            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"#,
        )
        .bind(model_name)
        .fetch_all(conn.as_mut())
        .await?;

        let mut versions = Vec::new();
        for (
            version,
            status_str,
            deployment_status_str,
            created_at,
            file_size,
            performance_metrics,
        ) in rows
        {
            versions.push(ModelVersion {
                version,
                status: self.parse_model_status(&status_str)?,
                deployment_status: self.parse_deployment_status(&deployment_status_str)?,
                created_at,
                file_size,
                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 tx = self.db.begin_transaction().await?;

        for dep in dependencies {
            let query = format!(
                r#"INSERT INTO ml_model_dependencies
                   (parent_model_id, dependency_model_id, dependency_type, weight)
                   VALUES ('{}', '{}', '{}', {})"#,
                parent_model_id,
                dep.model_id,
                dep.dependency_type.replace("'", "''"),
                dep.weight
            );
            tx.execute(&query).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 mut conn = self.db.acquire().await?;

        let rows = sqlx::query_as::<_, (Uuid, String, f64)>(
            r#"SELECT dependency_model_id, dependency_type, weight
               FROM ml_model_dependencies 
               WHERE parent_model_id = $1"#,
        )
        .bind(model_id)
        .fetch_all(conn.as_mut())
        .await?;

        let mut dependencies = Vec::new();
        for (dependency_model_id, dependency_type, weight) in rows {
            dependencies.push(ModelDependency {
                model_id: dependency_model_id,
                dependency_type,
                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::{Digest, Sha256};
        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 mut conn = self.db.acquire().await?;
        let count: i64 = sqlx::query_scalar(
            "SELECT COUNT(*) FROM ml_model_versions WHERE model_name = $1 AND version = $2",
        )
        .bind(model_name)
        .bind(version)
        .fetch_one(conn.as_mut())
        .await?;

        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 mut conn = self.db.acquire().await?;
        let _: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM ml_model_versions")
            .fetch_one(conn.as_mut())
            .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 record for lightweight operations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelRecord {
    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,
}