foxhunt/crates/ml-data/src/features.rs

//! Feature Engineering Repository
//!
//! Manages feature computation, versioning, lineage tracking, and real-time
//! feature serving for ML models in HFT trading systems with PostgreSQL integration.

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

use crate::{FeatureStoreConfig, MlDataError, Result};
use database::Database;
use sqlx::Row;

/// Feature repository for ML feature engineering and serving
#[derive(Clone)]
pub struct FeatureRepository {
    db: Database,
    config: FeatureStoreConfig,
}

impl FeatureRepository {
    pub async fn new(db: Database, config: FeatureStoreConfig) -> Result<Self> {
        let repo = Self { db, config };
        repo.initialize_schema().await?;
        Ok(repo)
    }

    /// Initialize database schema for feature management
    pub async fn initialize_schema(&self) -> Result<()> {
        // Initialize schema

        // Feature sets table
        self.db
            .execute(
                r#"
            CREATE TABLE IF NOT EXISTS ml_feature_sets (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                name VARCHAR NOT NULL,
                version INTEGER NOT NULL,
                description TEXT,
                created_at TIMESTAMPTZ DEFAULT NOW(),
                updated_at TIMESTAMPTZ DEFAULT NOW(),
                created_by VARCHAR NOT NULL,
                status feature_status DEFAULT 'draft',
                schema_definition JSONB NOT NULL,
                computation_config JSONB DEFAULT '{}',
                metadata JSONB DEFAULT '{}',
                UNIQUE(name, version)
            )
        "#,
            )
            .await?;

        // Feature definitions table
        self.db
            .execute(
                r#"
            CREATE TABLE IF NOT EXISTS ml_feature_definitions (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                feature_set_id UUID NOT NULL REFERENCES ml_feature_sets(id) ON DELETE CASCADE,
                name VARCHAR NOT NULL,
                data_type VARCHAR NOT NULL,
                description TEXT,
                computation_logic TEXT NOT NULL,
                dependencies TEXT[] DEFAULT '{}',
                transformation_type VARCHAR NOT NULL,
                default_value JSONB,
                validation_rules JSONB DEFAULT '{}',
                metadata JSONB DEFAULT '{}'
            )
        "#,
            )
            .await?;

        // Feature values table (for offline features)
        self.db
            .execute(
                r#"
            CREATE TABLE IF NOT EXISTS ml_feature_values (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                feature_set_id UUID NOT NULL REFERENCES ml_feature_sets(id) ON DELETE CASCADE,
                entity_id VARCHAR NOT NULL,
                timestamp TIMESTAMPTZ NOT NULL,
                features JSONB NOT NULL,
                version INTEGER NOT NULL,
                created_at TIMESTAMPTZ DEFAULT NOW(),
                expires_at TIMESTAMPTZ
            )
        "#,
            )
            .await?;

        // Feature lineage tracking
        self.db.execute(r#"
            CREATE TABLE IF NOT EXISTS ml_feature_lineage (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                downstream_feature_id UUID NOT NULL REFERENCES ml_feature_definitions(id) ON DELETE CASCADE,
                upstream_feature_id UUID REFERENCES ml_feature_definitions(id) ON DELETE CASCADE,
                upstream_data_source VARCHAR,
                transformation_type VARCHAR NOT NULL,
                dependency_type lineage_dependency_type NOT NULL,
                created_at TIMESTAMPTZ DEFAULT NOW(),
                metadata JSONB DEFAULT '{}'
            )
        "#).await?;

        // Feature computation jobs
        self.db
            .execute(
                r#"
            CREATE TABLE IF NOT EXISTS ml_feature_jobs (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                job_name VARCHAR NOT NULL,
                feature_set_id UUID NOT NULL REFERENCES ml_feature_sets(id) ON DELETE CASCADE,
                job_type job_type_enum NOT NULL,
                schedule_cron VARCHAR,
                started_at TIMESTAMPTZ,
                completed_at TIMESTAMPTZ,
                status job_status DEFAULT 'pending',
                records_processed BIGINT DEFAULT 0,
                error_message TEXT,
                configuration JSONB DEFAULT '{}',
                created_by VARCHAR NOT NULL
            )
        "#,
            )
            .await?;

        // Feature serving cache (for online features)
        self.db
            .execute(
                r#"
            CREATE TABLE IF NOT EXISTS ml_feature_cache (
                entity_id VARCHAR NOT NULL,
                feature_set_name VARCHAR NOT NULL,
                feature_set_version INTEGER NOT NULL,
                features JSONB NOT NULL,
                last_updated TIMESTAMPTZ NOT NULL,
                expires_at TIMESTAMPTZ NOT NULL,
                PRIMARY KEY (entity_id, feature_set_name, feature_set_version)
            )
        "#,
            )
            .await?;

        // Create enums
        self.db
            .execute(
                r#"
            DO $$ BEGIN
                CREATE TYPE feature_status AS ENUM ('draft', 'active', 'deprecated', 'archived');
            EXCEPTION
                WHEN duplicate_object THEN null;
            END $$;
        "#,
            )
            .await?;

        self.db.execute(r#"
            DO $$ BEGIN
                CREATE TYPE lineage_dependency_type AS ENUM ('direct', 'indirect', 'temporal', 'aggregation');
            EXCEPTION
                WHEN duplicate_object THEN null;
            END $$;
        "#).await?;

        self.db
            .execute(
                r#"
            DO $$ BEGIN
                CREATE TYPE job_type_enum AS ENUM ('batch', 'streaming', 'backfill', 'validation');
            EXCEPTION
                WHEN duplicate_object THEN null;
            END $$;
        "#,
            )
            .await?;

        self.db.execute(r#"
            DO $$ BEGIN
                CREATE TYPE job_status AS ENUM ('pending', 'running', 'completed', 'failed', 'cancelled');
            EXCEPTION
                WHEN duplicate_object THEN null;
            END $$;
        "#).await?;

        // Indexes for performance
        self.db.execute("CREATE INDEX IF NOT EXISTS idx_feature_sets_name_version ON ml_feature_sets(name, version)").await?;
        self.db.execute("CREATE INDEX IF NOT EXISTS idx_feature_values_entity_timestamp ON ml_feature_values(entity_id, timestamp DESC)").await?;
        self.db.execute("CREATE INDEX IF NOT EXISTS idx_feature_values_set_timestamp ON ml_feature_values(feature_set_id, timestamp DESC)").await?;
        self.db.execute("CREATE INDEX IF NOT EXISTS idx_feature_cache_lookup ON ml_feature_cache(entity_id, feature_set_name, feature_set_version)").await?;
        self.db.execute("CREATE INDEX IF NOT EXISTS idx_feature_cache_expires ON ml_feature_cache(expires_at)").await?;
        self.db.execute("CREATE INDEX IF NOT EXISTS idx_feature_jobs_status ON ml_feature_jobs(status, started_at DESC)").await?;

        Ok(())
    }

    /// Create a new feature set
    pub async fn create_feature_set(&self, request: CreateFeatureSetRequest) -> Result<FeatureSet> {
        // Validate feature set request
        self.validate_feature_set_request(&request).await?;

        // Check for version conflicts
        if self
            .feature_set_version_exists(&request.name, request.version)
            .await?
        {
            return Err(MlDataError::VersionConflict {
                message: format!(
                    "Feature set {} version {} already exists",
                    request.name, request.version
                ),
            });
        }

        let feature_set_id = Uuid::new_v4();

        // Clone needed values for the closure
        let name = request.name.clone();
        let version = request.version;
        let description = request.description.clone();
        let created_by = request.created_by.clone();
        let schema_definition = request.schema_definition.clone();
        let computation_config = request.computation_config.clone();
        let metadata = request.metadata.clone();
        let features = request.features.clone();

        // Execute transaction
        let mut tx = self.db.begin_transaction().await?;

        // Insert feature set record
        let description_sql = description
            .as_ref()
            .map(|s| format!("'{}'", s.replace("'", "''")))
            .unwrap_or("NULL".to_string());
        let query = format!(
            r#"INSERT INTO ml_feature_sets
               (id, name, version, description, created_by, schema_definition,
                computation_config, metadata)
               VALUES ('{}', '{}', {}, {}, '{}', '{}', '{}', '{}')"#,
            feature_set_id,
            name,
            version,
            description_sql,
            created_by,
            schema_definition.to_string().replace("'", "''"),
            computation_config.to_string().replace("'", "''"),
            metadata.to_string().replace("'", "''")
        );
        tx.execute(&query).await?;

        // Insert feature definitions
        let mut feature_defs = Vec::new();
        for feature_def in features {
            let feature_id = Uuid::new_v4();

            let description = feature_def
                .description
                .as_ref()
                .map(|s| format!("'{}'", s.replace("'", "''")))
                .unwrap_or("NULL".to_string());
            let default_value = feature_def
                .default_value
                .as_ref()
                .map(|v| format!("'{}'", v.to_string().replace("'", "''")))
                .unwrap_or("NULL".to_string());
            let dependencies_json =
                serde_json::to_string(&feature_def.dependencies).unwrap_or("[]".to_string());
            let query = format!(
                r#"INSERT INTO ml_feature_definitions
                   (id, feature_set_id, name, data_type, description, computation_logic,
                    dependencies, transformation_type, default_value, validation_rules, metadata)
                   VALUES ('{}', '{}', '{}', '{}', {}, '{}', '{}', '{}', {}, '{}', '{}')"#,
                feature_id,
                feature_set_id,
                feature_def.name,
                feature_def.data_type,
                description,
                feature_def.computation_logic,
                dependencies_json,
                feature_def.transformation_type,
                default_value,
                feature_def.validation_rules.to_string().replace("'", "''"),
                feature_def.metadata.to_string().replace("'", "''")
            );
            tx.execute(&query).await?;
            feature_defs.push(FeatureDefinition {
                id: feature_id,
                name: feature_def.name,
                data_type: feature_def.data_type,
                description: feature_def.description,
                computation_logic: feature_def.computation_logic,
                dependencies: feature_def.dependencies,
                transformation_type: feature_def.transformation_type,
                default_value: feature_def.default_value,
                validation_rules: feature_def.validation_rules,
                metadata: feature_def.metadata,
            });
        }

        tx.commit().await?;

        let feature_set = FeatureSet {
            id: feature_set_id,
            name: request.name,
            version: request.version,
            description: request.description,
            created_at: Utc::now(),
            updated_at: Utc::now(),
            created_by: request.created_by,
            status: FeatureSetStatus::Draft,
            schema_definition: request.schema_definition,
            computation_config: request.computation_config,
            metadata: request.metadata,
            features: feature_defs,
        };

        tracing::info!(
            "Created feature set: {} v{} with {} features",
            feature_set.name,
            feature_set.version,
            feature_set.features.len()
        );

        Ok(feature_set)
    }

    /// Compute and store feature values
    pub async fn compute_features(
        &self,
        feature_set_id: Uuid,
        entity_id: String,
        input_data: serde_json::Value,
        timestamp: Option<DateTime<Utc>>,
    ) -> Result<ComputedFeatures> {
        let timestamp = timestamp.unwrap_or_else(Utc::now);

        // Load feature set definition
        let feature_set = self.load_feature_set_by_id(feature_set_id).await?;

        // Compute features based on definitions
        let computed_values = self
            .execute_feature_computation(&feature_set, input_data)
            .await?;

        // Store computed features
        let mut conn = self.db.acquire().await?;
        sqlx::query(
            r#"INSERT INTO ml_feature_values 
               (feature_set_id, entity_id, timestamp, features, version, expires_at)
               VALUES ($1, $2, $3, $4, $5, $6)"#,
        )
        .bind(feature_set_id)
        .bind(&entity_id)
        .bind(timestamp)
        .bind(&computed_values)
        .bind(feature_set.version)
        .bind(self.calculate_expiry_time(timestamp))
        .execute(conn.as_mut())
        .await?;

        // Update serving cache if configured for online serving
        self.update_serving_cache(
            &entity_id,
            &feature_set.name,
            feature_set.version,
            &computed_values,
            timestamp,
        )
        .await?;

        Ok(ComputedFeatures {
            feature_set_id,
            entity_id,
            timestamp,
            features: computed_values,
            version: feature_set.version,
        })
    }

    /// Get features for online serving
    pub async fn get_online_features(
        &self,
        entity_id: &str,
        feature_set_name: &str,
        feature_set_version: Option<i32>,
    ) -> Result<Option<ServedFeatures>> {
        let mut conn = self.db.acquire().await?;

        // Using sqlx query builder pattern
        let result = if let Some(version) = feature_set_version {
            sqlx::query_as::<_, (serde_json::Value, DateTime<Utc>, DateTime<Utc>)>(
                r#"SELECT features, last_updated, expires_at
                   FROM ml_feature_cache 
                   WHERE entity_id = $1 AND feature_set_name = $2 AND feature_set_version = $3
                   AND expires_at > NOW()"#,
            )
            .bind(entity_id)
            .bind(feature_set_name)
            .bind(version)
            .fetch_optional(conn.as_mut())
            .await
        } else {
            sqlx::query_as::<_, (serde_json::Value, DateTime<Utc>, DateTime<Utc>)>(
                r#"SELECT features, last_updated, expires_at
                   FROM ml_feature_cache 
                   WHERE entity_id = $1 AND feature_set_name = $2
                   AND expires_at > NOW()
                   ORDER BY feature_set_version DESC
                   LIMIT 1"#,
            )
            .bind(entity_id)
            .bind(feature_set_name)
            .fetch_optional(conn.as_mut())
            .await
        };

        match result {
            Ok(Some((features, last_updated, expires_at))) => Ok(Some(ServedFeatures {
                entity_id: entity_id.to_string(),
                features,
                last_updated,
                expires_at,
            })),
            Ok(None) => Ok(None),
            Err(_) => Ok(None),
        }
    }

    /// Get historical feature values
    pub async fn get_historical_features(
        &self,
        feature_set_id: Uuid,
        entity_ids: Vec<String>,
        time_range: Option<(DateTime<Utc>, DateTime<Utc>)>,
        limit: Option<usize>,
    ) -> Result<Vec<HistoricalFeatures>> {
        let mut conn = self.db.acquire().await?;

        let _query = r#"
            SELECT entity_id, timestamp, features, version
            FROM ml_feature_values
            WHERE feature_set_id = $1
        "#
        .to_string();

        // Using sqlx query builder pattern
        let mut query_builder = sqlx::QueryBuilder::new(
            "SELECT feature_name, feature_value, computation_timestamp FROM ml_feature_vectors WHERE feature_set_id = "
        );
        query_builder.push_bind(feature_set_id);

        // Add entity filter
        if !entity_ids.is_empty() {
            query_builder.push(" AND entity_id = ANY(");
            query_builder.push_bind(&entity_ids);
            query_builder.push(")");
        }

        // Add time range filter
        if let Some((start, end)) = time_range {
            query_builder.push(" AND timestamp >= ");
            query_builder.push_bind(start);
            query_builder.push(" AND timestamp <= ");
            query_builder.push_bind(end);
        }

        query_builder.push(" ORDER BY timestamp DESC");

        // Add limit
        if let Some(limit_val) = limit {
            query_builder.push(" LIMIT ");
            query_builder.push_bind(limit_val as i64);
        }

        let query = query_builder.build();
        let rows = query.fetch_all(conn.as_mut()).await?;

        let mut results = Vec::new();
        for row in rows {
            results.push(HistoricalFeatures {
                entity_id: row.get("entity_id"),
                timestamp: row.get("timestamp"),
                features: row.get("features"),
                version: row.get("version"),
            });
        }

        Ok(results)
    }

    /// Track feature lineage
    pub async fn track_lineage(&self, lineage: FeatureLineage) -> Result<()> {
        let mut conn = self.db.acquire().await?;

        sqlx::query(
            r#"INSERT INTO ml_feature_lineage 
               (downstream_feature_id, upstream_feature_id, upstream_data_source,
                transformation_type, dependency_type, metadata)
               VALUES ($1, $2, $3, $4, $5, $6)"#,
        )
        .bind(lineage.downstream_feature_id)
        .bind(lineage.upstream_feature_id)
        .bind(&lineage.upstream_data_source)
        .bind(&lineage.transformation_type)
        .bind(lineage.dependency_type.to_string())
        .bind(&lineage.metadata)
        .execute(conn.as_mut())
        .await?;

        tracing::info!(
            "Tracked feature lineage for feature {}",
            lineage.downstream_feature_id
        );
        Ok(())
    }

    /// Start a feature computation job
    pub async fn start_feature_job(&self, request: StartFeatureJobRequest) -> Result<FeatureJob> {
        let mut conn = self.db.acquire().await?;
        let job_id = Uuid::new_v4();

        sqlx::query(
            r#"INSERT INTO ml_feature_jobs 
               (id, job_name, feature_set_id, job_type, schedule_cron, 
                started_at, configuration, created_by)
               VALUES ($1, $2, $3, $4, $5, $6, $7, $8)"#,
        )
        .bind(job_id)
        .bind(&request.job_name)
        .bind(request.feature_set_id)
        .bind(request.job_type.to_string())
        .bind(&request.schedule_cron)
        .bind(request.started_at)
        .bind(&request.configuration)
        .bind(&request.created_by)
        .execute(conn.as_mut())
        .await?;

        let job = FeatureJob {
            id: job_id,
            job_name: request.job_name,
            feature_set_id: request.feature_set_id,
            job_type: request.job_type,
            schedule_cron: request.schedule_cron,
            started_at: Some(request.started_at),
            completed_at: None,
            status: JobStatus::Running,
            records_processed: 0_i64,
            error_message: None,
            configuration: request.configuration,
            created_by: request.created_by,
        };

        tracing::info!("Started feature job: {}", job.job_name);
        Ok(job)
    }

    /// Execute feature computation logic
    async fn execute_feature_computation(
        &self,
        feature_set: &FeatureSet,
        input_data: serde_json::Value,
    ) -> Result<serde_json::Value> {
        let mut computed_features = serde_json::Map::new();

        for feature in &feature_set.features {
            let computed_value = self.compute_single_feature(feature, &input_data).await?;
            computed_features.insert(feature.name.clone(), computed_value);
        }

        Ok(serde_json::Value::Object(computed_features))
    }

    /// Compute a single feature value
    async fn compute_single_feature(
        &self,
        feature_def: &FeatureDefinition,
        input_data: &serde_json::Value,
    ) -> Result<serde_json::Value> {
        // This is a simplified implementation - in production, you'd have
        // a more sophisticated feature computation engine
        Ok(match feature_def.transformation_type.as_str() {
            "passthrough" => input_data
                .get(&feature_def.name)
                .cloned()
                .or_else(|| feature_def.default_value.clone())
                .unwrap_or(serde_json::Value::Null),
            "normalize" => {
                if let Some(value) = input_data.get(&feature_def.name).and_then(|v| v.as_f64()) {
                    // Simple min-max normalization (would be configurable)
                    serde_json::json!((value - 0.0) / (1.0 - 0.0))
                } else {
                    feature_def
                        .default_value
                        .clone()
                        .unwrap_or(serde_json::Value::Null)
                }
            },
            "log_transform" => {
                if let Some(value) = input_data.get(&feature_def.name).and_then(|v| v.as_f64()) {
                    if value > 0.0 {
                        serde_json::json!(value.ln())
                    } else {
                        serde_json::json!(0.0)
                    }
                } else {
                    feature_def
                        .default_value
                        .clone()
                        .unwrap_or(serde_json::Value::Null)
                }
            },
            _ => feature_def
                .default_value
                .clone()
                .unwrap_or(serde_json::Value::Null),
        })
    }

    /// Update serving cache for online features
    async fn update_serving_cache(
        &self,
        entity_id: &str,
        feature_set_name: &str,
        feature_set_version: i32,
        features: &serde_json::Value,
        timestamp: DateTime<Utc>,
    ) -> Result<()> {
        let mut conn = self.db.acquire().await?;
        let expires_at = self.calculate_expiry_time(timestamp);

        sqlx::query(
            r#"INSERT INTO ml_feature_cache 
               (entity_id, feature_set_name, feature_set_version, features, last_updated, expires_at)
               VALUES ($1, $2, $3, $4, $5, $6)
               ON CONFLICT (entity_id, feature_set_name, feature_set_version)
               DO UPDATE SET features = EXCLUDED.features, last_updated = EXCLUDED.last_updated,
                           expires_at = EXCLUDED.expires_at"#
        )
        .bind(entity_id)
        .bind(feature_set_name)
        .bind(feature_set_version)
        .bind(features)
        .bind(timestamp)
        .bind(expires_at)
        .execute(conn.as_mut())
        .await?;

        Ok(())
    }

    /// Calculate feature expiry time based on TTL configuration
    fn calculate_expiry_time(&self, timestamp: DateTime<Utc>) -> DateTime<Utc> {
        timestamp + Duration::seconds(self.config.ttl_seconds as i64)
    }

    /// Load feature set by ID
    async fn load_feature_set_by_id(&self, feature_set_id: Uuid) -> Result<FeatureSet> {
        let mut conn = self.db.acquire().await?;

        // Load feature set
        let set_row = sqlx::query_as::<
            _,
            (
                String,
                i32,
                Option<String>,
                DateTime<Utc>,
                DateTime<Utc>,
                String,
                String,
                serde_json::Value,
                serde_json::Value,
                serde_json::Value,
            ),
        >(
            r#"SELECT name, version, description, created_at, updated_at, created_by,
                      status, schema_definition, computation_config, metadata
               FROM ml_feature_sets WHERE id = $1"#,
        )
        .bind(feature_set_id)
        .fetch_one(conn.as_mut())
        .await
        .map_err(|_| MlDataError::NotFound {
            resource_type: "FeatureSet".to_string(),
            id: feature_set_id.to_string(),
        })?;

        // Load feature definitions
        let feature_rows = sqlx::query_as::<
            _,
            (
                Uuid,
                String,
                String,
                Option<String>,
                String,
                Vec<String>,
                String,
                Option<serde_json::Value>,
                serde_json::Value,
                serde_json::Value,
            ),
        >(
            r#"SELECT id, name, data_type, description, computation_logic, dependencies,
                      transformation_type, default_value, validation_rules, metadata
               FROM ml_feature_definitions WHERE feature_set_id = $1"#,
        )
        .bind(feature_set_id)
        .fetch_all(conn.as_mut())
        .await?;

        let mut features = Vec::new();
        for (
            id,
            name,
            data_type,
            description,
            computation_logic,
            dependencies,
            transformation_type,
            default_value,
            validation_rules,
            metadata,
        ) in feature_rows
        {
            features.push(FeatureDefinition {
                id,
                name,
                data_type,
                description,
                computation_logic,
                dependencies,
                transformation_type,
                default_value,
                validation_rules,
                metadata,
            });
        }
        let (
            name,
            version,
            description,
            created_at,
            updated_at,
            created_by,
            status,
            schema_definition,
            computation_config,
            metadata,
        ) = set_row;

        Ok(FeatureSet {
            id: feature_set_id,
            name,
            version,
            description,
            created_at,
            updated_at,
            created_by,
            status: self.parse_feature_set_status(&status)?,
            schema_definition,
            computation_config,
            metadata,
            features,
        })
    }

    /// Validate feature set request
    async fn validate_feature_set_request(&self, request: &CreateFeatureSetRequest) -> Result<()> {
        if request.name.trim().is_empty() {
            return Err(MlDataError::Validation {
                message: "Feature set name cannot be empty".to_string(),
            });
        }

        if request.version < 1 {
            return Err(MlDataError::Validation {
                message: "Feature set version must be >= 1".to_string(),
            });
        }

        if request.features.is_empty() {
            return Err(MlDataError::Validation {
                message: "Feature set must contain at least one feature".to_string(),
            });
        }

        Ok(())
    }

    /// Check if feature set version exists
    async fn feature_set_version_exists(&self, name: &str, version: i32) -> Result<bool> {
        let mut conn = self.db.acquire().await?;
        let count: i64 = sqlx::query_scalar(
            "SELECT COUNT(*) FROM ml_feature_sets WHERE name = $1 AND version = $2",
        )
        .bind(name)
        .bind(version)
        .fetch_one(conn.as_mut())
        .await?;

        Ok(count > 0)
    }

    /// Parse feature set status from database
    fn parse_feature_set_status(&self, status_str: &str) -> Result<FeatureSetStatus> {
        match status_str {
            "draft" => Ok(FeatureSetStatus::Draft),
            "active" => Ok(FeatureSetStatus::Active),
            "deprecated" => Ok(FeatureSetStatus::Deprecated),
            "archived" => Ok(FeatureSetStatus::Archived),
            _ => Err(MlDataError::Validation {
                message: format!("Invalid feature set status: {}", status_str),
            }),
        }
    }

    /// Health check for feature 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_feature_sets")
            .fetch_one(conn.as_mut())
            .await?;
        Ok(true)
    }
}

/// Request to create a new feature set
#[derive(Debug)]
pub struct CreateFeatureSetRequest {
    pub name: String,
    pub version: i32,
    pub description: Option<String>,
    pub created_by: String,
    pub schema_definition: serde_json::Value,
    pub computation_config: serde_json::Value,
    pub metadata: serde_json::Value,
    pub features: Vec<CreateFeatureDefinitionRequest>,
}

/// Request to create a feature definition
#[derive(Debug, Clone)]
pub struct CreateFeatureDefinitionRequest {
    pub name: String,
    pub data_type: String,
    pub description: Option<String>,
    pub computation_logic: String,
    pub dependencies: Vec<String>,
    pub transformation_type: String,
    pub default_value: Option<serde_json::Value>,
    pub validation_rules: serde_json::Value,
    pub metadata: serde_json::Value,
}

/// Feature set representation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeatureSet {
    pub id: Uuid,
    pub name: String,
    pub version: i32,
    pub description: Option<String>,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
    pub created_by: String,
    pub status: FeatureSetStatus,
    pub schema_definition: serde_json::Value,
    pub computation_config: serde_json::Value,
    pub metadata: serde_json::Value,
    pub features: Vec<FeatureDefinition>,
}

/// Feature set status
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum FeatureSetStatus {
    Draft,
    Active,
    Deprecated,
    Archived,
}

/// Feature definition
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeatureDefinition {
    pub id: Uuid,
    pub name: String,
    pub data_type: String,
    pub description: Option<String>,
    pub computation_logic: String,
    pub dependencies: Vec<String>,
    pub transformation_type: String,
    pub default_value: Option<serde_json::Value>,
    pub validation_rules: serde_json::Value,
    pub metadata: serde_json::Value,
}

/// Computed feature values
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ComputedFeatures {
    pub feature_set_id: Uuid,
    pub entity_id: String,
    pub timestamp: DateTime<Utc>,
    pub features: serde_json::Value,
    pub version: i32,
}

/// Features served for online inference
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServedFeatures {
    pub entity_id: String,
    pub features: serde_json::Value,
    pub last_updated: DateTime<Utc>,
    pub expires_at: DateTime<Utc>,
}

/// Historical feature values
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HistoricalFeatures {
    pub entity_id: String,
    pub timestamp: DateTime<Utc>,
    pub features: serde_json::Value,
    pub version: i32,
}

/// Feature lineage tracking
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeatureLineage {
    pub downstream_feature_id: Uuid,
    pub upstream_feature_id: Option<Uuid>,
    pub upstream_data_source: Option<String>,
    pub transformation_type: String,
    pub dependency_type: LineageDependencyType,
    pub metadata: serde_json::Value,
}

/// Feature lineage dependency types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum LineageDependencyType {
    Direct,
    Indirect,
    Temporal,
    Aggregation,
}

impl ToString for LineageDependencyType {
    fn to_string(&self) -> String {
        match self {
            LineageDependencyType::Direct => "direct".to_string(),
            LineageDependencyType::Indirect => "indirect".to_string(),
            LineageDependencyType::Temporal => "temporal".to_string(),
            LineageDependencyType::Aggregation => "aggregation".to_string(),
        }
    }
}

/// Feature computation job request
#[derive(Debug)]
pub struct StartFeatureJobRequest {
    pub job_name: String,
    pub feature_set_id: Uuid,
    pub job_type: JobType,
    pub schedule_cron: Option<String>,
    pub started_at: DateTime<Utc>,
    pub configuration: serde_json::Value,
    pub created_by: String,
}

/// Feature computation job
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeatureJob {
    pub id: Uuid,
    pub job_name: String,
    pub feature_set_id: Uuid,
    pub job_type: JobType,
    pub schedule_cron: Option<String>,
    pub started_at: Option<DateTime<Utc>>,
    pub completed_at: Option<DateTime<Utc>>,
    pub status: JobStatus,
    pub records_processed: i64,
    pub error_message: Option<String>,
    pub configuration: serde_json::Value,
    pub created_by: String,
}

/// Job types for feature computation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum JobType {
    Batch,
    Streaming,
    Backfill,
    Validation,
}

impl ToString for JobType {
    fn to_string(&self) -> String {
        match self {
            JobType::Batch => "batch".to_string(),
            JobType::Streaming => "streaming".to_string(),
            JobType::Backfill => "backfill".to_string(),
            JobType::Validation => "validation".to_string(),
        }
    }
}

/// Job execution status
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum JobStatus {
    Pending,
    Running,
    Completed,
    Failed,
    Cancelled,
}

/// Feature version information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeatureVersion {
    pub name: String,
    pub version: i32,
    pub created_at: DateTime<Utc>,
    pub status: FeatureSetStatus,
    pub feature_count: usize,
}