//! Job Tracker Module
//!
//! Implements batch job hierarchy tracking and progress aggregation for ML training jobs.
//! Supports parent batch jobs with multiple child training jobs and weighted progress calculation.

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use sqlx::PgPool;
use tracing::{debug, info};
use uuid::Uuid;

/// Training job status
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum JobStatus {
    Pending,
    Running,
    Completed,
    Failed,
    Stopped,
    Paused,
}

impl JobStatus {
    /// Convert to database string representation
    pub fn to_db_string(&self) -> &'static str {
        match self {
            JobStatus::Pending => "Pending",
            JobStatus::Running => "Running",
            JobStatus::Completed => "Completed",
            JobStatus::Failed => "Failed",
            JobStatus::Stopped => "Stopped",
            JobStatus::Paused => "Paused",
        }
    }

    /// Parse from database string
    pub fn from_db_string(s: &str) -> Result<Self> {
        match s {
            "Pending" => Ok(JobStatus::Pending),
            "Running" => Ok(JobStatus::Running),
            "Completed" => Ok(JobStatus::Completed),
            "Failed" => Ok(JobStatus::Failed),
            "Stopped" => Ok(JobStatus::Stopped),
            "Paused" => Ok(JobStatus::Paused),
            _ => Err(anyhow::anyhow!("Invalid job status: {}", s)),
        }
    }

    /// Validate state transition
    pub fn is_valid_transition(&self, new_status: &JobStatus) -> bool {
        use JobStatus::*;
        
        match (self, new_status) {
            // From Pending
            (Pending, Running) => true,
            (Pending, Stopped) => true,
            (Pending, Failed) => true,
            
            // From Running
            (Running, Completed) => true,
            (Running, Failed) => true,
            (Running, Paused) => true,
            (Running, Stopped) => true,
            
            // From Paused
            (Paused, Running) => true,
            (Paused, Stopped) => true,
            (Paused, Failed) => true,
            
            // Terminal states cannot transition
            (Completed, _) => false,
            (Failed, _) => false,
            (Stopped, _) => false,
            
            // Same state is allowed (idempotent)
            (a, b) if a == b => true,
            
            // All other transitions are invalid
            _ => false,
        }
    }
}

/// Batch job summary with aggregated metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BatchJobSummary {
    pub batch_id: Uuid,
    pub total_jobs: u32,
    pub pending_jobs: u32,
    pub running_jobs: u32,
    pub completed_jobs: u32,
    pub failed_jobs: u32,
    pub overall_progress: f64,  // 0.0-100.0
    pub status: JobStatus,
}

/// Job progress information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JobProgress {
    pub job_id: Uuid,
    pub current_epoch: u32,
    pub total_epochs: u32,
    pub progress_pct: f64,  // 0.0-100.0
}

/// Child job record from database
#[derive(Debug, Clone, sqlx::FromRow)]
struct ChildJob {
    model_weight: f64,
    status: String,
    progress_pct: f64,
}

/// Job tracker for managing batch job hierarchies and progress aggregation
pub struct JobTracker {
    db_pool: PgPool,
}

impl JobTracker {
    /// Create a new job tracker
    pub fn new(db_pool: PgPool) -> Self {
        Self { db_pool }
    }

    /// Update job status with validation
    pub async fn update_job_status(
        &self,
        job_id: Uuid,
        status: JobStatus,
    ) -> Result<()> {
        debug!("Updating job {} to status {:?}", job_id, status);

        // Get current status
        let current_status: String = sqlx::query_scalar(
            "SELECT status FROM child_jobs WHERE id = $1"
        )
        .bind(job_id)
        .fetch_one(&self.db_pool)
        .await
        .context("Failed to fetch current job status")?;

        let current = JobStatus::from_db_string(&current_status)?;

        // Validate state transition
        if !current.is_valid_transition(&status) {
            return Err(anyhow::anyhow!(
                "Invalid state transition from {:?} to {:?} for job {}",
                current,
                status,
                job_id
            ));
        }

        // Update status - use raw query to avoid CASE statement type inference issues
        let status_str = status.to_db_string();
        
        sqlx::query(
            r"
            UPDATE child_jobs
            SET status = $1,
                started_at = CASE
                    WHEN $1 = 'Running' AND started_at IS NULL THEN NOW()
                    ELSE started_at
                END,
                completed_at = CASE
                    WHEN $1 IN ('Completed', 'Failed', 'Stopped') THEN NOW()
                    ELSE completed_at
                END
            WHERE id = $2
            "
        )
        .bind(status_str)
        .bind(job_id)
        .execute(&self.db_pool)
        .await
        .context("Failed to update job status")?;

        info!("Updated job {} to status {:?}", job_id, status);

        Ok(())
    }

    /// Update job progress
    pub async fn update_job_progress(
        &self,
        job_id: Uuid,
        progress: JobProgress,
    ) -> Result<()> {
        debug!(
            "Updating job {} progress: {}/{}",
            job_id, progress.current_epoch, progress.total_epochs
        );

        sqlx::query(
            r"
            UPDATE child_jobs
            SET current_epoch = $1,
                total_epochs = $2,
                progress_pct = $3
            WHERE id = $4
            "
        )
        .bind(progress.current_epoch as i32)
        .bind(progress.total_epochs as i32)
        .bind(progress.progress_pct)
        .bind(job_id)
        .execute(&self.db_pool)
        .await
        .context("Failed to update job progress")?;

        info!(
            "Updated job {} progress to {:.1}%",
            job_id, progress.progress_pct
        );

        Ok(())
    }

    /// Get batch job summary with aggregated metrics
    pub async fn get_batch_summary(
        &self,
        batch_id: Uuid,
    ) -> Result<BatchJobSummary> {
        debug!("Fetching batch summary for batch {}", batch_id);

        // Verify batch exists
        let batch_exists: bool = sqlx::query_scalar(
            "SELECT EXISTS(SELECT 1 FROM batch_jobs WHERE id = $1)"
        )
        .bind(batch_id)
        .fetch_one(&self.db_pool)
        .await
        .context("Failed to verify batch existence")?;

        if !batch_exists {
            return Err(anyhow::anyhow!("Batch job {} not found", batch_id));
        }

        // Fetch all child jobs for this batch
        let child_jobs: Vec<ChildJob> = sqlx::query_as::<_, ChildJob>(
            r#"
            SELECT model_weight, status, progress_pct
            FROM child_jobs
            WHERE batch_id = $1
            "#,
        )
        .bind(batch_id)
        .fetch_all(&self.db_pool)
        .await
        .context("Failed to fetch child jobs")?;

        // Count jobs by status
        let total_jobs = child_jobs.len() as u32;
        let pending_jobs = child_jobs.iter()
            .filter(|j| j.status == "Pending")
            .count() as u32;
        let running_jobs = child_jobs.iter()
            .filter(|j| j.status == "Running")
            .count() as u32;
        let completed_jobs = child_jobs.iter()
            .filter(|j| j.status == "Completed")
            .count() as u32;
        let failed_jobs = child_jobs.iter()
            .filter(|j| j.status == "Failed")
            .count() as u32;

        // Calculate weighted progress
        let overall_progress = self.calculate_weighted_progress(&child_jobs);

        // Determine batch status
        let status = self.determine_batch_status(
            total_jobs,
            pending_jobs,
            running_jobs,
            completed_jobs,
            failed_jobs,
        );

        let summary = BatchJobSummary {
            batch_id,
            total_jobs,
            pending_jobs,
            running_jobs,
            completed_jobs,
            failed_jobs,
            overall_progress,
            status,
        };

        debug!(
            "Batch {} summary: {}/{} completed, {:.1}% progress",
            batch_id, completed_jobs, total_jobs, overall_progress
        );

        Ok(summary)
    }

    /// Calculate weighted progress across all child jobs
    ///
    /// Weights: DQN:10%, PPO:30%, MAMBA-2:40%, TFT:20%
    /// Formula: Sum(job_progress * model_weight) for all jobs
    fn calculate_weighted_progress(&self, child_jobs: &[ChildJob]) -> f64 {
        if child_jobs.is_empty() {
            return 0.0;
        }

        let weighted_sum: f64 = child_jobs
            .iter()
            .map(|job| job.progress_pct * job.model_weight)
            .sum();

        // Progress is already in percentage (0.0-100.0)
        weighted_sum
    }

    /// Determine batch status based on child job counts
    fn determine_batch_status(
        &self,
        total_jobs: u32,
        pending_jobs: u32,
        running_jobs: u32,
        completed_jobs: u32,
        failed_jobs: u32,
    ) -> JobStatus {
        // If any jobs are running, batch is running
        if running_jobs > 0 {
            return JobStatus::Running;
        }

        // If all jobs are completed, batch is completed
        if total_jobs > 0 && completed_jobs == total_jobs {
            return JobStatus::Completed;
        }

        // If there are failures and no running/pending jobs, batch failed
        if failed_jobs > 0 && running_jobs == 0 && pending_jobs == 0 {
            return JobStatus::Failed;
        }

        // Default to pending
        JobStatus::Pending
    }

    /// Create a new batch job
    pub async fn create_batch_job(
        &self,
        batch_id: Uuid,
        name: String,
        description: String,
    ) -> Result<()> {
        debug!("Creating batch job {} ({})", batch_id, name);

        sqlx::query(
            r"
            INSERT INTO batch_jobs (id, name, description, status)
            VALUES ($1, $2, $3, $4)
            "
        )
        .bind(batch_id)
        .bind(name)
        .bind(description)
        .bind("Pending")
        .execute(&self.db_pool)
        .await
        .context("Failed to create batch job")?;

        info!("Created batch job {}", batch_id);

        Ok(())
    }

    /// Create a child job
    pub async fn create_child_job(
        &self,
        job_id: Uuid,
        batch_id: Uuid,
        model_type: String,
        model_weight: f64,
        total_epochs: u32,
    ) -> Result<()> {
        debug!(
            "Creating child job {} for batch {} ({})",
            job_id, batch_id, model_type
        );

        sqlx::query(
            r"
            INSERT INTO child_jobs (
                id, batch_id, model_type, model_weight,
                status, progress_pct, current_epoch, total_epochs
            )
            VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
            "
        )
        .bind(job_id)
        .bind(batch_id)
        .bind(model_type.clone())
        .bind(model_weight)
        .bind("Pending")
        .bind(0.0)
        .bind(0i32)
        .bind(total_epochs as i32)
        .execute(&self.db_pool)
        .await
        .context("Failed to create child job")?;

        info!(
            "Created child job {} for batch {} ({})",
            job_id, batch_id, model_type
        );

        Ok(())
    }
}

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

    #[test]
    fn test_job_status_valid_transitions() {
        assert!(JobStatus::Pending.is_valid_transition(&JobStatus::Running));
        assert!(JobStatus::Running.is_valid_transition(&JobStatus::Completed));
        assert!(JobStatus::Running.is_valid_transition(&JobStatus::Failed));
        assert!(JobStatus::Paused.is_valid_transition(&JobStatus::Running));
    }

    #[test]
    fn test_job_status_invalid_transitions() {
        assert!(!JobStatus::Completed.is_valid_transition(&JobStatus::Running));
        assert!(!JobStatus::Failed.is_valid_transition(&JobStatus::Running));
        assert!(!JobStatus::Pending.is_valid_transition(&JobStatus::Completed));
    }

    #[test]
    fn test_job_status_same_state() {
        assert!(JobStatus::Pending.is_valid_transition(&JobStatus::Pending));
        assert!(JobStatus::Running.is_valid_transition(&JobStatus::Running));
    }

    #[tokio::test]
    async fn test_calculate_weighted_progress_empty() {
        let pool = PgPool::connect_lazy("postgresql://dummy").unwrap();
        let tracker = JobTracker::new(pool);
        let jobs: Vec<ChildJob> = vec![];
        assert_eq!(tracker.calculate_weighted_progress(&jobs), 0.0);
    }

    #[tokio::test]
    async fn test_calculate_weighted_progress_standard_weights() {
        let pool = PgPool::connect_lazy("postgresql://dummy").unwrap();
        let tracker = JobTracker::new(pool);
        
        let jobs = vec![
            ChildJob {
                model_weight: 0.10,
                status: "Completed".to_string(),
                progress_pct: 100.0,
            },
            ChildJob {
                model_weight: 0.30,
                status: "Running".to_string(),
                progress_pct: 50.0,
            },
            ChildJob {
                model_weight: 0.40,
                status: "Running".to_string(),
                progress_pct: 25.0,
            },
            ChildJob {
                model_weight: 0.20,
                status: "Pending".to_string(),
                progress_pct: 0.0,
            },
        ];
        
        // Expected: (100*0.10) + (50*0.30) + (25*0.40) + (0*0.20) = 10 + 15 + 10 + 0 = 35.0
        let progress = tracker.calculate_weighted_progress(&jobs);
        assert!((progress - 35.0).abs() < 0.01);
    }

    #[tokio::test]
    async fn test_determine_batch_status_all_pending() {
        let pool = PgPool::connect_lazy("postgresql://dummy").unwrap();
        let tracker = JobTracker::new(pool);

        let status = tracker.determine_batch_status(4, 4, 0, 0, 0);
        assert_eq!(status, JobStatus::Pending);
    }

    #[tokio::test]
    async fn test_determine_batch_status_running() {
        let pool = PgPool::connect_lazy("postgresql://dummy").unwrap();
        let tracker = JobTracker::new(pool);

        let status = tracker.determine_batch_status(4, 2, 2, 0, 0);
        assert_eq!(status, JobStatus::Running);
    }

    #[tokio::test]
    async fn test_determine_batch_status_completed() {
        let pool = PgPool::connect_lazy("postgresql://dummy").unwrap();
        let tracker = JobTracker::new(pool);

        let status = tracker.determine_batch_status(4, 0, 0, 4, 0);
        assert_eq!(status, JobStatus::Completed);
    }

    #[tokio::test]
    async fn test_determine_batch_status_failed() {
        let pool = PgPool::connect_lazy("postgresql://dummy").unwrap();
        let tracker = JobTracker::new(pool);

        let status = tracker.determine_batch_status(4, 0, 0, 2, 2);
        assert_eq!(status, JobStatus::Failed);
    }
}
