//! Repository Layer for ML Training Service
//!
//! This module defines the repository abstraction for ML training data persistence,
//! implementing the repository pattern to decouple business logic from database concerns.

use std::collections::HashMap;
use async_trait::async_trait;
use anyhow::Result;
use uuid::Uuid;

use crate::database::{DatabaseManager, TrainingJobRecord};
use crate::orchestrator::JobStatus;

/// Repository trait for ML training job data operations
#[async_trait]
pub trait MlDataRepository: Send + Sync {
    /// Insert a new training job
    async fn create_training_job(&self, job_record: &TrainingJobRecord) -> Result<()>;

    /// Update an existing training job
    async fn update_training_job(&self, job_record: &TrainingJobRecord) -> Result<()>;

    /// Get a training job by ID
    async fn find_training_job(&self, job_id: Uuid) -> Result<Option<TrainingJobRecord>>;

    /// List training jobs with filtering and pagination
    async fn list_training_jobs(
        &self,
        status_filter: Option<&str>,
        model_type_filter: Option<&str>,
        limit: Option<i64>,
        offset: Option<i64>,
    ) -> Result<Vec<TrainingJobRecord>>;

    /// Get training job count with optional filters
    async fn count_training_jobs(
        &self,
        status_filter: Option<&str>,
        model_type_filter: Option<&str>,
    ) -> Result<i64>;

    /// Insert training metrics for an epoch
    async fn save_training_metrics(
        &self,
        job_id: Uuid,
        epoch: i32,
        train_loss: Option<f64>,
        validation_loss: Option<f64>,
        metrics: &HashMap<String, f64>,
    ) -> Result<()>;

    /// Get training metrics for a job
    async fn get_training_metrics(
        &self,
        job_id: Uuid,
    ) -> Result<Vec<(i32, f32, f32, HashMap<String, f64>)>>;

    /// Delete a training job and its metrics
    async fn delete_training_job(&self, job_id: Uuid) -> Result<bool>;

    /// Health check for repository connectivity
    async fn health_check(&self) -> Result<()>;
}

/// Postgre`SQL` implementation of MlDataRepository
pub struct PostgresMlDataRepository {
    database: std::sync::Arc<DatabaseManager>,
}

impl PostgresMlDataRepository {
    /// Create a new Postgre`SQL` repository
    pub fn new(database: std::sync::Arc<DatabaseManager>) -> Self {
        Self { database }
    }
}

#[async_trait]
impl MlDataRepository for PostgresMlDataRepository {
    async fn create_training_job(&self, job_record: &TrainingJobRecord) -> Result<()> {
        self.database.insert_training_job(job_record).await
    }

    async fn update_training_job(&self, job_record: &TrainingJobRecord) -> Result<()> {
        self.database.update_training_job(job_record).await
    }

    async fn find_training_job(&self, job_id: Uuid) -> Result<Option<TrainingJobRecord>> {
        self.database.get_training_job(job_id).await
    }

    async fn list_training_jobs(
        &self,
        status_filter: Option<&str>,
        model_type_filter: Option<&str>,
        limit: Option<i64>,
        offset: Option<i64>,
    ) -> Result<Vec<TrainingJobRecord>> {
        self.database
            .list_training_jobs(status_filter, model_type_filter, limit, offset)
            .await
    }

    async fn count_training_jobs(
        &self,
        status_filter: Option<&str>,
        model_type_filter: Option<&str>,
    ) -> Result<i64> {
        self.database
            .count_training_jobs(status_filter, model_type_filter)
            .await
    }

    async fn save_training_metrics(
        &self,
        job_id: Uuid,
        epoch: i32,
        train_loss: Option<f64>,
        validation_loss: Option<f64>,
        metrics: &HashMap<String, f64>,
    ) -> Result<()> {
        self.database
            .insert_training_metrics(job_id, epoch, train_loss, validation_loss, metrics)
            .await
    }

    async fn get_training_metrics(
        &self,
        job_id: Uuid,
    ) -> Result<Vec<(i32, f32, f32, HashMap<String, f64>)>> {
        self.database.get_training_metrics(job_id).await
    }

    async fn delete_training_job(&self, job_id: Uuid) -> Result<bool> {
        self.database.delete_training_job(job_id).await
    }

    async fn health_check(&self) -> Result<()> {
        self.database.health_check().await
    }
}

/// Mock implementation for testing
#[cfg(test)]
pub struct MockMlDataRepository {
    jobs: std::sync::Arc<tokio::sync::RwLock<HashMap<Uuid, TrainingJobRecord>>>,
    metrics: std::sync::Arc<tokio::sync::RwLock<HashMap<Uuid, Vec<(i32, f32, f32, HashMap<String, f64>)>>>>,
}

#[cfg(test)]
impl MockMlDataRepository {
    pub fn new() -> Self {
        Self {
            jobs: std::sync::Arc::new(tokio::sync::RwLock::new(HashMap::new())),
            metrics: std::sync::Arc::new(tokio::sync::RwLock::new(HashMap::new())),
        }
    }
}

#[cfg(test)]
#[async_trait]
impl MlDataRepository for MockMlDataRepository {
    async fn create_training_job(&self, job_record: &TrainingJobRecord) -> Result<()> {
        let mut jobs = self.jobs.write().await;
        jobs.insert(job_record.id, job_record.clone());
        Ok(())
    }

    async fn update_training_job(&self, job_record: &TrainingJobRecord) -> Result<()> {
        let mut jobs = self.jobs.write().await;
        jobs.insert(job_record.id, job_record.clone());
        Ok(())
    }

    async fn find_training_job(&self, job_id: Uuid) -> Result<Option<TrainingJobRecord>> {
        let jobs = self.jobs.read().await;
        Ok(jobs.get(&job_id).cloned())
    }

    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 jobs = self.jobs.read().await;
        let filtered_jobs: Vec<TrainingJobRecord> = jobs
            .values()
            .filter(|job| {
                status_filter
                    .map(|status| job.status == status)
                    .unwrap_or(true)
            })
            .cloned()
            .collect();
        Ok(filtered_jobs)
    }

    async fn count_training_jobs(
        &self,
        _status_filter: Option<&str>,
        _model_type_filter: Option<&str>,
    ) -> Result<i64> {
        let jobs = self.jobs.read().await;
        Ok(jobs.len() as i64)
    }

    async fn save_training_metrics(
        &self,
        job_id: Uuid,
        epoch: i32,
        train_loss: Option<f64>,
        validation_loss: Option<f64>,
        metrics: &HashMap<String, f64>,
    ) -> Result<()> {
        let mut job_metrics = self.metrics.write().await;
        let entry = job_metrics.entry(job_id).or_insert_with(Vec::new);
        entry.push((
            epoch,
            train_loss.unwrap_or(0.0) as f32,
            validation_loss.unwrap_or(0.0) as f32,
            metrics.clone(),
        ));
        Ok(())
    }

    async fn get_training_metrics(
        &self,
        job_id: Uuid,
    ) -> Result<Vec<(i32, f32, f32, HashMap<String, f64>)>> {
        let job_metrics = self.metrics.read().await;
        Ok(job_metrics.get(&job_id).cloned().unwrap_or_default())
    }

    async fn delete_training_job(&self, job_id: Uuid) -> Result<bool> {
        let mut jobs = self.jobs.write().await;
        let mut job_metrics = self.metrics.write().await;
        jobs.remove(&job_id);
        job_metrics.remove(&job_id);
        Ok(true)
    }

    async fn health_check(&self) -> Result<()> {
        Ok(())
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;
    use chrono::Utc;

    #[tokio::test]
    async fn test_mock_repository() {
        let repo = MockMlDataRepository::new();
        
        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,
        };

        // Test create
        repo.create_training_job(&job_record).await.unwrap();

        // Test find
        let found = repo.find_training_job(job_record.id).await.unwrap();
        assert!(found.is_some());
        assert_eq!(found.unwrap().id, job_record.id);

        // Test list
        let jobs = repo.list_training_jobs(None, None, None, None).await.unwrap();
        assert_eq!(jobs.len(), 1);

        // Test count
        let count = repo.count_training_jobs(None, None).await.unwrap();
        assert_eq!(count, 1);

        // Test health check
        repo.health_check().await.unwrap();
    }
}