foxhunt/services/ml_training_service/src/database.rs

//! Database Management for ML Training Service
//!
//! This module handles PostgreSQL database operations for storing training job metadata,
//! configurations, and results.

use std::collections::HashMap;

use anyhow::{Context, Result};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sqlx::Row;
use tracing::{debug, info};
use uuid::Uuid;

use crate::orchestrator::{JobStatus, TrainingJob};
use common::database::DatabasePool;
use config::database::DatabaseConfig;

/// Database manager for training job persistence
pub struct DatabaseManager {
    db_pool: DatabasePool,
}

/// Training job record for database storage
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingJobRecord {
    pub id: Uuid,
    pub model_type: String,
    pub status: String,
    pub config_json: String,
    pub created_at: DateTime<Utc>,
    pub started_at: Option<DateTime<Utc>>,
    pub completed_at: Option<DateTime<Utc>>,
    pub description: String,
    pub tags_json: String,
    pub progress_percentage: f32,
    pub current_epoch: i32,
    pub total_epochs: i32,
    pub metrics_json: String,
    pub error_message: Option<String>,
    pub model_artifact_path: Option<String>,
}

impl TrainingJobRecord {
    /// Convert from TrainingJob to database record
    pub fn from_training_job(job: &TrainingJob) -> Self {
        Self {
            id: job.id,
            model_type: job.model_type.clone(),
            status: format!("{:?}", job.status),
            config_json: serde_json::to_string(&job.config).unwrap_or_default(),
            created_at: job.created_at,
            started_at: job.started_at,
            completed_at: job.completed_at,
            description: job.description.clone(),
            tags_json: serde_json::to_string(&job.tags).unwrap_or_default(),
            progress_percentage: job.progress_percentage,
            current_epoch: i32::try_from(job.current_epoch).unwrap_or(0),
            total_epochs: i32::try_from(job.total_epochs).unwrap_or(0),
            metrics_json: serde_json::to_string(&job.metrics).unwrap_or_default(),
            error_message: job.error_message.clone(),
            model_artifact_path: job.model_artifact_path.clone(),
        }
    }

    /// Convert from database record to TrainingJob
    pub fn to_training_job(&self) -> Result<TrainingJob> {
        let status = match self.status.as_str() {
            "Pending" => JobStatus::Pending,
            "Running" => JobStatus::Running,
            "Completed" => JobStatus::Completed,
            "Failed" => JobStatus::Failed,
            "Stopped" => JobStatus::Stopped,
            "Paused" => JobStatus::Paused,
            _ => JobStatus::Pending,
        };

        let config =
            serde_json::from_str(&self.config_json).context("Failed to deserialize config")?;

        let tags: HashMap<String, String> =
            serde_json::from_str(&self.tags_json).unwrap_or_default();

        let metrics: HashMap<String, f64> =
            serde_json::from_str(&self.metrics_json).unwrap_or_default();

        Ok(TrainingJob {
            id: self.id,
            model_type: self.model_type.clone(),
            status,
            config,
            created_at: self.created_at,
            started_at: self.started_at,
            completed_at: self.completed_at,
            description: self.description.clone(),
            tags,
            progress_percentage: self.progress_percentage,
            current_epoch: u32::try_from(self.current_epoch).unwrap_or(0),
            total_epochs: u32::try_from(self.total_epochs).unwrap_or(0),
            metrics,
            error_message: self.error_message.clone(),
            model_artifact_path: self.model_artifact_path.clone(),
        })
    }
}

impl DatabaseManager {
    /// Create a new database manager
    pub async fn new(config: &DatabaseConfig) -> Result<Self> {
        Self::new_with_migrations(config, true).await
    }

    /// Create a new database manager with optional migrations
    pub async fn new_with_migrations(
        config: &DatabaseConfig,
        run_migrations: bool,
    ) -> Result<Self> {
        info!(
            "Connecting to database: {}",
            config.url.replace([':', '@'], "*")
        );

        // Convert config::DatabaseConfig to common::database::LocalDatabaseConfig using From trait
        let common_config: common::database::LocalDatabaseConfig = config.clone().into();

        let db_pool = DatabasePool::new(common_config)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to create database pool: {}", e))?;

        info!("Database connection pool established with HFT optimizations");

        let manager = Self { db_pool };

        // Run migrations if requested
        if run_migrations {
            manager.run_migrations().await?;
        }

        Ok(manager)
    }

    /// Get a clone of the underlying `PgPool` for use by other components (e.g. `JobSpawner`).
    ///
    /// `PgPool` is an `Arc`-wrapped handle, so cloning is cheap.
    pub fn pg_pool(&self) -> sqlx::PgPool {
        self.db_pool.pool().clone()
    }

    /// Run database migrations
    pub async fn run_migrations(&self) -> Result<()> {
        info!("Running database migrations");

        // Use advisory lock to prevent concurrent migrations
        sqlx::query("SELECT pg_advisory_lock(123456789)")
            .execute(self.db_pool.pool())
            .await?;

        // Create training jobs table
        sqlx::query(
            r#"
            CREATE TABLE IF NOT EXISTS training_jobs (
                id UUID PRIMARY KEY,
                model_type VARCHAR NOT NULL,
                status VARCHAR NOT NULL,
                config_json TEXT NOT NULL,
                created_at TIMESTAMPTZ NOT NULL,
                started_at TIMESTAMPTZ,
                completed_at TIMESTAMPTZ,
                description TEXT NOT NULL,
                tags_json TEXT NOT NULL DEFAULT '{}',
                progress_percentage REAL NOT NULL DEFAULT 0.0,
                current_epoch INTEGER NOT NULL DEFAULT 0,
                total_epochs INTEGER NOT NULL DEFAULT 0,
                metrics_json TEXT NOT NULL DEFAULT '{}',
                error_message TEXT,
                model_artifact_path TEXT,
                
                -- Indexes for common queries
                CONSTRAINT training_jobs_status_check CHECK (
                    status IN ('Pending', 'Running', 'Completed', 'Failed', 'Stopped', 'Paused')
                )
            )
            "#,
        )
        .execute(self.db_pool.pool())
        .await
        .context("Failed to create training_jobs table")?;

        // Create indexes
        sqlx::query("CREATE INDEX IF NOT EXISTS idx_training_jobs_status ON training_jobs(status)")
            .execute(self.db_pool.pool())
            .await?;

        sqlx::query(
            "CREATE INDEX IF NOT EXISTS idx_training_jobs_model_type ON training_jobs(model_type)",
        )
        .execute(self.db_pool.pool())
        .await?;

        sqlx::query("CREATE INDEX IF NOT EXISTS idx_training_jobs_created_at ON training_jobs(created_at DESC)")
            .execute(self.db_pool.pool())
            .await?;

        // Create training metrics table for detailed tracking
        sqlx::query(
            r#"
            CREATE TABLE IF NOT EXISTS training_metrics (
                id UUID DEFAULT gen_random_uuid() PRIMARY KEY,
                job_id UUID NOT NULL REFERENCES training_jobs(id) ON DELETE CASCADE,
                epoch INTEGER NOT NULL,
                timestamp TIMESTAMPTZ NOT NULL,
                train_loss REAL,
                validation_loss REAL,
                metrics_json TEXT NOT NULL DEFAULT '{}',
                
                UNIQUE(job_id, epoch)
            )
            "#,
        )
        .execute(self.db_pool.pool())
        .await
        .context("Failed to create training_metrics table")?;

        sqlx::query("CREATE INDEX IF NOT EXISTS idx_training_metrics_job_id ON training_metrics(job_id, epoch)")
            .execute(self.db_pool.pool())
            .await?;

        // Release advisory lock
        sqlx::query("SELECT pg_advisory_unlock(123456789)")
            .execute(self.db_pool.pool())
            .await?;

        info!("Database migrations completed successfully");
        Ok(())
    }

    /// Insert a new training job
    pub async fn insert_training_job(&self, job: &TrainingJobRecord) -> Result<()> {
        sqlx::query(
            r#"
            INSERT INTO training_jobs (
                id, model_type, status, config_json, created_at, started_at, completed_at,
                description, tags_json, progress_percentage, current_epoch, total_epochs,
                metrics_json, error_message, model_artifact_path
            ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)
            "#,
        )
        .bind(job.id)
        .bind(&job.model_type)
        .bind(&job.status)
        .bind(&job.config_json)
        .bind(job.created_at)
        .bind(job.started_at)
        .bind(job.completed_at)
        .bind(&job.description)
        .bind(&job.tags_json)
        .bind(job.progress_percentage)
        .bind(job.current_epoch)
        .bind(job.total_epochs)
        .bind(&job.metrics_json)
        .bind(&job.error_message)
        .bind(&job.model_artifact_path)
        .execute(self.db_pool.pool())
        .await
        .context("Failed to insert training job")?;

        debug!("Inserted training job {}", job.id);
        Ok(())
    }

    /// Update an existing training job
    pub async fn update_training_job(&self, job: &TrainingJobRecord) -> Result<()> {
        sqlx::query(
            r#"
            UPDATE training_jobs SET
                status = $2,
                started_at = $3,
                completed_at = $4,
                progress_percentage = $5,
                current_epoch = $6,
                total_epochs = $7,
                metrics_json = $8,
                error_message = $9,
                model_artifact_path = $10
            WHERE id = $1
            "#,
        )
        .bind(job.id)
        .bind(&job.status)
        .bind(job.started_at)
        .bind(job.completed_at)
        .bind(job.progress_percentage)
        .bind(job.current_epoch)
        .bind(job.total_epochs)
        .bind(&job.metrics_json)
        .bind(&job.error_message)
        .bind(&job.model_artifact_path)
        .execute(self.db_pool.pool())
        .await
        .context("Failed to update training job")?;

        debug!("Updated training job {}", job.id);
        Ok(())
    }

    /// Get a training job by ID
    pub async fn get_training_job(&self, job_id: Uuid) -> Result<Option<TrainingJobRecord>> {
        let row = sqlx::query(
            r#"
            SELECT id, model_type, status, config_json, created_at, started_at, completed_at,
                   description, tags_json, progress_percentage, current_epoch, total_epochs,
                   metrics_json, error_message, model_artifact_path
            FROM training_jobs
            WHERE id = $1
            "#,
        )
        .bind(job_id)
        .fetch_optional(self.db_pool.pool())
        .await
        .context("Failed to fetch training job")?;

        if let Some(row) = row {
            let record = TrainingJobRecord {
                id: row.get("id"),
                model_type: row.get("model_type"),
                status: row.get("status"),
                config_json: row.get("config_json"),
                created_at: row.get("created_at"),
                started_at: row.get("started_at"),
                completed_at: row.get("completed_at"),
                description: row.get("description"),
                tags_json: row.get("tags_json"),
                progress_percentage: row.get("progress_percentage"),
                current_epoch: row.get("current_epoch"),
                total_epochs: row.get("total_epochs"),
                metrics_json: row.get("metrics_json"),
                error_message: row.get("error_message"),
                model_artifact_path: row.get("model_artifact_path"),
            };
            Ok(Some(record))
        } else {
            Ok(None)
        }
    }

    /// List training jobs with filtering and pagination
    pub async fn list_training_jobs(
        &self,
        status_filter: Option<&str>,
        model_type_filter: Option<&str>,
        limit: Option<i64>,
        offset: Option<i64>,
    ) -> Result<Vec<TrainingJobRecord>> {
        let mut query = String::from(
            r#"
            SELECT id, model_type, status, config_json, created_at, started_at, completed_at,
                   description, tags_json, progress_percentage, current_epoch, total_epochs,
                   metrics_json, error_message, model_artifact_path
            FROM training_jobs
            WHERE 1=1
            "#,
        );

        let mut bind_count = 0;

        if status_filter.is_some() {
            bind_count += 1;
            query.push_str(&format!(" AND status = ${}", bind_count));
        }

        if model_type_filter.is_some() {
            bind_count += 1;
            query.push_str(&format!(" AND model_type = ${}", bind_count));
        }

        query.push_str(" ORDER BY created_at DESC");

        if limit.is_some() {
            bind_count += 1;
            query.push_str(&format!(" LIMIT ${}", bind_count));
        }

        if offset.is_some() {
            bind_count += 1;
            query.push_str(&format!(" OFFSET ${}", bind_count));
        }

        let mut sql_query = sqlx::query(&query);

        if let Some(status) = status_filter {
            sql_query = sql_query.bind(status);
        }

        if let Some(model_type) = model_type_filter {
            sql_query = sql_query.bind(model_type);
        }

        if let Some(limit) = limit {
            sql_query = sql_query.bind(limit);
        }

        if let Some(offset) = offset {
            sql_query = sql_query.bind(offset);
        }

        let rows = sql_query
            .fetch_all(self.db_pool.pool())
            .await
            .context("Failed to fetch training jobs")?;

        let mut jobs = Vec::new();
        for row in rows {
            let record = TrainingJobRecord {
                id: row.get("id"),
                model_type: row.get("model_type"),
                status: row.get("status"),
                config_json: row.get("config_json"),
                created_at: row.get("created_at"),
                started_at: row.get("started_at"),
                completed_at: row.get("completed_at"),
                description: row.get("description"),
                tags_json: row.get("tags_json"),
                progress_percentage: row.get("progress_percentage"),
                current_epoch: row.get("current_epoch"),
                total_epochs: row.get("total_epochs"),
                metrics_json: row.get("metrics_json"),
                error_message: row.get("error_message"),
                model_artifact_path: row.get("model_artifact_path"),
            };
            jobs.push(record);
        }

        Ok(jobs)
    }

    /// Get training job count with optional filters
    pub async fn count_training_jobs(
        &self,
        status_filter: Option<&str>,
        model_type_filter: Option<&str>,
    ) -> Result<i64> {
        let mut query = String::from("SELECT COUNT(*) FROM training_jobs WHERE 1=1");
        let mut bind_count = 0;

        if status_filter.is_some() {
            bind_count += 1;
            query.push_str(&format!(" AND status = ${}", bind_count));
        }

        if model_type_filter.is_some() {
            bind_count += 1;
            query.push_str(&format!(" AND model_type = ${}", bind_count));
        }

        let mut sql_query = sqlx::query_scalar(&query);

        if let Some(status) = status_filter {
            sql_query = sql_query.bind(status);
        }

        if let Some(model_type) = model_type_filter {
            sql_query = sql_query.bind(model_type);
        }

        let count: i64 = sql_query
            .fetch_one(self.db_pool.pool())
            .await
            .context("Failed to count training jobs")?;

        Ok(count)
    }

    /// Insert training metrics for an epoch
    pub async fn insert_training_metrics(
        &self,
        job_id: Uuid,
        epoch: i32,
        train_loss: Option<f64>,
        validation_loss: Option<f64>,
        metrics: &HashMap<String, f64>,
    ) -> Result<()> {
        let metrics_json = serde_json::to_string(metrics).context("Failed to serialize metrics")?;

        sqlx::query(
            r#"
            INSERT INTO training_metrics (job_id, epoch, timestamp, train_loss, validation_loss, metrics_json)
            VALUES ($1, $2, $3, $4, $5, $6)
            ON CONFLICT (job_id, epoch) DO UPDATE SET
                timestamp = EXCLUDED.timestamp,
                train_loss = EXCLUDED.train_loss,
                validation_loss = EXCLUDED.validation_loss,
                metrics_json = EXCLUDED.metrics_json
            "#
        )
        .bind(job_id)
        .bind(epoch)
        .bind(Utc::now())
        .bind(train_loss.map(|x| x as f32))
        .bind(validation_loss.map(|x| x as f32))
        .bind(metrics_json)
        .execute(self.db_pool.pool())
        .await
        .context("Failed to insert training metrics")?;

        debug!(
            "Inserted training metrics for job {} epoch {}",
            job_id, epoch
        );
        Ok(())
    }

    /// Get training metrics for a job
    pub async fn get_training_metrics(
        &self,
        job_id: Uuid,
    ) -> Result<Vec<(i32, f32, f32, HashMap<String, f64>)>> {
        let rows = sqlx::query(
            r#"
            SELECT epoch, train_loss, validation_loss, metrics_json
            FROM training_metrics
            WHERE job_id = $1
            ORDER BY epoch
            "#,
        )
        .bind(job_id)
        .fetch_all(self.db_pool.pool())
        .await
        .context("Failed to fetch training metrics")?;

        let mut metrics = Vec::new();
        for row in rows {
            let epoch: i32 = row.get("epoch");
            let train_loss: Option<f32> = row.get("train_loss");
            let validation_loss: Option<f32> = row.get("validation_loss");
            let metrics_json: String = row.get("metrics_json");

            let parsed_metrics: HashMap<String, f64> =
                serde_json::from_str(&metrics_json).unwrap_or_default();

            metrics.push((
                epoch,
                train_loss.unwrap_or(0.0),
                validation_loss.unwrap_or(0.0),
                parsed_metrics,
            ));
        }

        Ok(metrics)
    }

    /// Delete a training job and its metrics
    pub async fn delete_training_job(&self, job_id: Uuid) -> Result<bool> {
        let result = sqlx::query("DELETE FROM training_jobs WHERE id = $1")
            .bind(job_id)
            .execute(self.db_pool.pool())
            .await
            .context("Failed to delete training job")?;

        Ok(result.rows_affected() > 0)
    }

    /// Health check for database connectivity
    pub async fn health_check(&self) -> Result<()> {
        sqlx::query("SELECT 1")
            .execute(self.db_pool.pool())
            .await
            .context("Database health check failed")?;

        Ok(())
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;
    use config::database::DatabaseConfig;
    // Note: These tests require a running PostgreSQL database
    // In a CI environment, you would use a test database

    async fn setup_test_db() -> Result<DatabaseManager> {
        let config = DatabaseConfig {
            url: "postgresql://test:test@localhost:5432/test_ml_training".to_string(),
            max_connections: 5,
            min_connections: 1,
            connect_timeout: std::time::Duration::from_secs(10),
            query_timeout: std::time::Duration::from_secs(30),
            enable_query_logging: true,
            application_name: Some("ml_training_service_test".to_string()),
            pool: config::database::PoolConfig::default(),
            transaction: config::database::TransactionConfig::default(),
        };

        DatabaseManager::new(&config).await
    }

    #[tokio::test]
    #[ignore = "Requires database setup"]
    async fn test_database_migrations() {
        let db = setup_test_db()
            .await
            .expect("Failed to setup test database");

        // Migrations should have run automatically
        assert!(db.health_check().await.is_ok());
    }

    #[tokio::test]
    #[ignore = "Requires database setup"]
    async fn test_insert_and_get_job() {
        let db = setup_test_db()
            .await
            .expect("Failed to setup test database");

        let job_record = TrainingJobRecord {
            id: Uuid::new_v4(),
            model_type: "TLOB".to_string(),
            status: "Pending".to_string(),
            config_json: "{}".to_string(),
            created_at: Utc::now(),
            started_at: None,
            completed_at: None,
            description: "Test job".to_string(),
            tags_json: "{}".to_string(),
            progress_percentage: 0.0,
            current_epoch: 0,
            total_epochs: 100,
            metrics_json: "{}".to_string(),
            error_message: None,
            model_artifact_path: None,
        };

        // Insert job
        db.insert_training_job(&job_record)
            .await
            .expect("Failed to insert job");

        // Get job
        let retrieved = db
            .get_training_job(job_record.id)
            .await
            .expect("Failed to get job");
        assert!(retrieved.is_some());

        let retrieved = retrieved.unwrap();
        assert_eq!(retrieved.id, job_record.id);
        assert_eq!(retrieved.model_type, job_record.model_type);
        assert_eq!(retrieved.status, job_record.status);
    }
}