//! Training Error Recovery and Edge Case Tests for ML Training Service
//!
//! Tests covering:
//! 1. Training pipeline errors (NaN, OOM, divergence)
//! 2. Checkpoint save/load failures
//! 3. GPU resource exhaustion
//! 4. Metric recording under errors
//! 5. Job cancellation during training
//! 6. Concurrent resource contention

use std::collections::HashMap;
use std::sync::Arc;

use ml::training_pipeline::{
    FinancialValidationConfig, ModelArchitectureConfig, PerformanceConfig,
    ProductionTrainingConfig, TrainingHyperparameters,
};
use ml_training_service::{
    checkpoint_manager::{CheckpointManager, RetentionPolicy},
    gpu_resource_manager::{GPUAllocationError, GPUResourceManager},
    orchestrator::{JobStatus, TrainingJob},
    training_metrics,
};
use sha2::Digest;
use sqlx::PgPool;
use uuid::Uuid;

/// Helper to create minimal training config
fn create_minimal_config() -> ProductionTrainingConfig {
    ProductionTrainingConfig {
        model_config: ModelArchitectureConfig {
            input_dim: 10,
            output_dim: 1,
            hidden_dims: vec![32],
            dropout_rate: 0.0,
            activation: "relu".to_string(),
            batch_norm: false,
            residual_connections: false,
        },
        training_params: TrainingHyperparameters {
            learning_rate: 0.001,
            batch_size: 8,
            max_epochs: 2,
            patience: 1,
            validation_split: 0.2,
            l2_regularization: 0.0,
            lr_decay_factor: 0.1,
            lr_decay_patience: 2,
        },
        safety_config: ml::safety::MLSafetyConfig::default(),
        gradient_config: ml::safety::GradientSafetyConfig::default(),
        financial_config: FinancialValidationConfig {
            max_prediction_multiple: 2.0,
            min_prediction_confidence: 0.5,
            validate_position_sizing: true,
            max_position_fraction: 0.25,
            min_sharpe_threshold: 0.5,
        },
        performance_config: PerformanceConfig {
            device_preference: "cpu".to_string(),
            max_memory_bytes: 1024 * 1024 * 1024, // 1GB
            num_workers: 1,
            gradient_accumulation_steps: 1,
        },
    }
}

/// Helper to setup test database
async fn setup_test_db() -> PgPool {
    let database_url = std::env::var("DATABASE_URL").unwrap_or_else(|_| {
        "postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt".to_string()
    });

    PgPool::connect(&database_url)
        .await
        .expect("Failed to connect to test database")
}

// ============================================================================
// Test 1: Checkpoint Save Failures
// ============================================================================

#[tokio::test]
async fn test_checkpoint_manager_handles_corrupted_checksum() {
    let pool = setup_test_db().await;

    let manager = CheckpointManager::new(pool.clone(), RetentionPolicy::default())
        .await
        .expect("Failed to create checkpoint manager");

    // Create a checkpoint with known checksum
    let checkpoint_data = b"test checkpoint data";
    let mut hasher = sha2::Sha256::new();
    hasher.update(checkpoint_data);
    let correct_checksum = format!("{:x}", hasher.finalize());

    // Use unique version to avoid conflicts
    let test_version = format!(
        "1.0.{}",
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs()
    );

    // Register checkpoint
    let metadata = ml::checkpoint::CheckpointMetadata {
        checkpoint_id: Uuid::new_v4().to_string(),
        model_type: ml::ModelType::DQN,
        model_name: "test_corrupted".to_string(),
        version: test_version.clone(),
        created_at: chrono::Utc::now(),
        epoch: Some(10),
        step: Some(1000),
        loss: Some(0.1),
        accuracy: Some(0.9),
        hyperparameters: HashMap::new(),
        metrics: HashMap::new(),
        architecture: HashMap::new(),
        format: ml::checkpoint::CheckpointFormat::Binary,
        compression: ml::checkpoint::CompressionType::LZ4,
        file_size: checkpoint_data.len() as u64,
        compressed_size: None,
        checksum: correct_checksum.clone(),
        tags: vec![],
        custom_metadata: HashMap::new(),
        signature: None,
        signature_algorithm: "none".to_string(),
        signing_key_id: "test".to_string(),
        signed_at: None,
    };

    let checkpoint_id = manager
        .register_checkpoint(metadata.clone())
        .await
        .expect("Failed to register checkpoint");

    // Test 1: Valid checksum passes
    let result = manager
        .validate_checksum(&checkpoint_id, checkpoint_data)
        .await;
    assert!(result.is_ok(), "Valid checksum should pass");

    // Test 2: Corrupted data fails
    let corrupted_data = b"corrupted checkpoint data";
    let result = manager
        .validate_checksum(&checkpoint_id, corrupted_data)
        .await;
    assert!(result.is_err(), "Corrupted checksum should fail");

    if let Err(e) = result {
        let error_msg = format!("{}", e);
        assert!(
            error_msg.contains("Checksum mismatch"),
            "Error should mention checksum mismatch, got: {}",
            error_msg
        );
    }

    // Cleanup
    let _ = sqlx::query("DELETE FROM ml_model_versions WHERE metadata->>'test_model_name' = $1")
        .bind("test_corrupted")
        .execute(&pool)
        .await;
}

// ============================================================================
// Test 2: GPU Resource Exhaustion
// ============================================================================

#[tokio::test]
async fn test_gpu_manager_rejects_insufficient_memory() {
    let manager = GPUResourceManager::new(vec![0])
        .await
        .expect("Failed to create GPU manager");

    let job_id = Uuid::new_v4();

    // Try to acquire GPU with impossibly large memory requirement
    let result = manager
        .acquire_gpu_with_memory_requirement(
            job_id,
            0,
            1_000_000_000, // 1 TB - impossible
        )
        .await;

    assert!(
        result.is_err(),
        "Should reject impossible memory requirement"
    );

    if let Err(GPUAllocationError::InsufficientMemory {
        gpu_id,
        required_mb,
        available_mb,
    }) = result
    {
        assert_eq!(gpu_id, 0);
        assert_eq!(required_mb, 1_000_000_000);
        assert!(available_mb < required_mb);
    } else {
        panic!("Expected InsufficientMemory error, got: {:?}", result);
    }
}

#[tokio::test]
async fn test_gpu_manager_handles_concurrent_allocation() {
    let manager = Arc::new(
        GPUResourceManager::new(vec![0])
            .await
            .expect("Failed to create GPU manager"),
    );

    let job1 = Uuid::new_v4();
    let job2 = Uuid::new_v4();

    // First job acquires GPU
    let lock1 = manager
        .acquire_gpu(job1, 0)
        .await
        .expect("First job should acquire GPU");

    assert_eq!(lock1.gpu_id(), 0);
    assert_eq!(lock1.job_id(), job1);

    // Second job should fail to acquire same GPU
    let result = manager.acquire_gpu(job2, 0).await;
    assert!(
        result.is_err(),
        "Second job should fail to acquire locked GPU"
    );

    if let Err(GPUAllocationError::GPUAlreadyLocked {
        gpu_id,
        current_job_id,
    }) = result
    {
        assert_eq!(gpu_id, 0);
        assert_eq!(current_job_id, job1);
    } else {
        panic!("Expected GPUAlreadyLocked error, got: {:?}", result);
    }

    // Release GPU from first job
    drop(lock1);

    // Give async drop time to execute
    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

    // Second job should now be able to acquire
    let lock2 = manager
        .acquire_gpu(job2, 0)
        .await
        .expect("Second job should acquire GPU after release");

    assert_eq!(lock2.gpu_id(), 0);
    assert_eq!(lock2.job_id(), job2);
}

#[tokio::test]
async fn test_gpu_manager_prevents_wrong_job_release() {
    let manager = GPUResourceManager::new(vec![0])
        .await
        .expect("Failed to create GPU manager");

    let job1 = Uuid::new_v4();
    let job2 = Uuid::new_v4();

    // Job 1 acquires GPU
    let _lock = manager
        .acquire_gpu(job1, 0)
        .await
        .expect("Job 1 should acquire GPU");

    // Job 2 tries to release GPU locked by Job 1
    let result = manager.release_gpu(0, job2).await;

    assert!(
        result.is_err(),
        "Should not allow releasing GPU locked by different job"
    );

    if let Err(GPUAllocationError::CannotReleaseLockedByDifferentJob {
        gpu_id,
        locked_job_id,
        requested_job_id,
    }) = result
    {
        assert_eq!(gpu_id, 0);
        assert_eq!(locked_job_id, job1);
        assert_eq!(requested_job_id, job2);
    } else {
        panic!(
            "Expected CannotReleaseLockedByDifferentJob error, got: {:?}",
            result
        );
    }
}

// ============================================================================
// Test 3: Training Metrics Under Errors
// ============================================================================

#[test]
fn test_training_metrics_records_nan_detection() {
    training_metrics::init_metrics();

    // Record NaN in different tensor types
    training_metrics::record_nan_detection("dqn", "job-123", "loss");
    training_metrics::record_nan_detection("dqn", "job-123", "gradient");
    training_metrics::record_nan_detection("dqn", "job-123", "activation");

    // Metrics should be recorded without panic
    // In production, Prometheus would track these
}

#[test]
fn test_training_metrics_records_checkpoint_failures() {
    training_metrics::init_metrics();

    // Record successful checkpoint
    training_metrics::record_checkpoint_save(
        "mamba",
        "job-456",
        true,
        5.2,
        100 * 1024 * 1024,
        None,
    );

    // Record failed checkpoint with different error types
    training_metrics::record_checkpoint_save("tft", "job-789", false, 0.0, 0, Some("disk_full"));

    training_metrics::record_checkpoint_save(
        "ppo",
        "job-101",
        false,
        0.0,
        0,
        Some("permission_denied"),
    );

    // Metrics should be recorded without panic
}

#[test]
fn test_training_metrics_records_gpu_metrics() {
    training_metrics::init_metrics();

    // Record GPU metrics
    training_metrics::record_gpu_metrics(
        "0",
        85.5,            // Utilization 85.5%
        6_500_000_000.0, // 6.5 GB used
        8_000_000_000.0, // 8 GB total
        75.0,            // 75°C temperature
    );

    // Record overheating GPU
    training_metrics::record_gpu_metrics(
        "1",
        100.0,           // 100% utilization
        7_800_000_000.0, // 7.8 GB used (near limit)
        8_000_000_000.0, // 8 GB total
        92.0,            // 92°C (hot!)
    );

    // Metrics should be recorded without panic
}

// ============================================================================
// Test 4: Job Lifecycle Edge Cases
// ============================================================================

#[test]
fn test_training_job_tracks_progress() {
    let config = create_minimal_config();
    let mut job = TrainingJob::new(
        "test_model".to_string(),
        config,
        "Test job".to_string(),
        HashMap::new(),
    );

    // Initial state
    assert_eq!(job.status, JobStatus::Pending);
    assert_eq!(job.progress_percentage, 0.0);

    // Simulate training progress
    job.status = JobStatus::Running;
    job.started_at = Some(chrono::Utc::now());
    job.current_epoch = 5;
    job.total_epochs = 10;
    job.progress_percentage = 50.0;

    assert_eq!(job.status, JobStatus::Running);
    assert_eq!(job.current_epoch, 5);
    assert_eq!(job.progress_percentage, 50.0);
    assert!(job.started_at.is_some());

    // Simulate completion
    job.status = JobStatus::Completed;
    job.completed_at = Some(chrono::Utc::now());
    job.progress_percentage = 100.0;
    job.metrics.insert("final_loss".to_string(), 0.05);

    assert_eq!(job.status, JobStatus::Completed);
    assert_eq!(job.progress_percentage, 100.0);
    assert!(job.completed_at.is_some());
    assert_eq!(job.metrics.get("final_loss"), Some(&0.05));
}

#[test]
fn test_training_job_tracks_failure() {
    let config = create_minimal_config();
    let mut job = TrainingJob::new(
        "failing_model".to_string(),
        config,
        "Test failure".to_string(),
        HashMap::new(),
    );

    // Simulate training failure
    job.status = JobStatus::Running;
    job.started_at = Some(chrono::Utc::now());
    job.current_epoch = 3;
    job.total_epochs = 10;

    // Failure occurs
    job.status = JobStatus::Failed;
    job.completed_at = Some(chrono::Utc::now());
    job.error_message = Some("NaN detected in loss".to_string());

    assert_eq!(job.status, JobStatus::Failed);
    assert!(job.error_message.is_some());
    assert!(job.error_message.unwrap().contains("NaN"));
    assert_eq!(job.current_epoch, 3);
    assert!(job.progress_percentage < 100.0);
}

// ============================================================================
// Test 5: Checkpoint Manager Concurrent Operations
// ============================================================================

#[tokio::test]
async fn test_checkpoint_manager_handles_concurrent_registrations() {
    let pool = setup_test_db().await;
    let manager = Arc::new(
        CheckpointManager::new(pool.clone(), RetentionPolicy::default())
            .await
            .expect("Failed to create checkpoint manager"),
    );

    // Create multiple checkpoints concurrently
    let mut handles = vec![];

    for i in 0..5 {
        let manager_clone = Arc::clone(&manager);
        let handle = tokio::spawn(async move {
            let unique_version = format!(
                "1.0.{}",
                std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap()
                    .as_nanos()
                    + i as u128
            );
            let metadata = ml::checkpoint::CheckpointMetadata {
                checkpoint_id: Uuid::new_v4().to_string(),
                model_type: ml::ModelType::DQN,
                model_name: format!("concurrent_test_{}", i),
                version: unique_version,
                created_at: chrono::Utc::now(),
                epoch: Some(10),
                step: Some(1000),
                loss: Some(0.1),
                accuracy: Some(0.9),
                hyperparameters: HashMap::new(),
                metrics: HashMap::from([("sharpe_ratio".to_string(), 1.5)]),
                architecture: HashMap::new(),
                format: ml::checkpoint::CheckpointFormat::Binary,
                compression: ml::checkpoint::CompressionType::LZ4,
                file_size: 1024,
                compressed_size: None,
                checksum: "0".repeat(64),
                tags: vec![],
                custom_metadata: HashMap::new(),
                signature: None,
                signature_algorithm: "none".to_string(),
                signing_key_id: "test".to_string(),
                signed_at: None,
            };

            manager_clone.register_checkpoint(metadata).await
        });

        handles.push(handle);
    }

    // Wait for all registrations
    let results: Vec<_> = futures::future::join_all(handles).await;

    // All should succeed
    for result in results {
        assert!(result.is_ok(), "Concurrent registration should succeed");
        let registration_result = result.unwrap();
        assert!(
            registration_result.is_ok(),
            "Checkpoint registration should succeed"
        );
    }

    // Cleanup
    for i in 0..5 {
        let model_name = format!("concurrent_test_{}", i);
        let _ =
            sqlx::query("DELETE FROM ml_model_versions WHERE metadata->>'test_model_name' = $1")
                .bind(&model_name)
                .execute(&pool)
                .await;
    }
}

// ============================================================================
// Test 6: Semantic Version Validation Edge Cases
// ============================================================================

#[tokio::test]
async fn test_checkpoint_manager_validates_semantic_versions() {
    let pool = setup_test_db().await;
    let manager = CheckpointManager::new(pool, RetentionPolicy::default())
        .await
        .expect("Failed to create checkpoint manager");

    // Valid versions
    let valid_versions = vec![
        "0.0.1",
        "1.0.0",
        "1.2.3",
        "10.20.30",
        "1.0.0-alpha",
        "1.0.0-beta.1",
        "1.0.0+build123",
        "1.0.0-rc1+build456",
    ];

    for version in valid_versions {
        let result = manager.validate_version(version).await;
        assert!(result.is_ok(), "Version '{}' should be valid", version);
    }

    // Invalid versions
    let invalid_versions = vec![
        "1",       // Missing minor/patch
        "1.0",     // Missing patch
        "v1.0.0",  // Leading 'v'
        "1.0.0.0", // Too many components
        "1.a.0",   // Non-numeric
        "a.b.c",   // All non-numeric
        "",        // Empty
        " 1.0.0 ", // Whitespace
    ];

    for version in invalid_versions {
        let result = manager.validate_version(version).await;
        assert!(result.is_err(), "Version '{}' should be invalid", version);
    }
}

// ============================================================================
// Test 7: GPU Statistics and Monitoring
// ============================================================================

#[tokio::test]
async fn test_gpu_manager_provides_accurate_statistics() {
    let manager = GPUResourceManager::new(vec![0, 1, 2])
        .await
        .expect("Failed to create GPU manager");

    // Initial statistics
    let stats = manager.get_statistics().await;
    assert_eq!(stats.total_gpus, 3);
    assert_eq!(stats.available_gpus, 3);
    assert_eq!(stats.locked_gpus, 0);
    assert_eq!(stats.active_jobs, 0);

    // Acquire some GPUs
    let job1 = Uuid::new_v4();
    let job2 = Uuid::new_v4();

    let _lock1 = manager
        .acquire_gpu(job1, 0)
        .await
        .expect("Should acquire GPU 0");
    let _lock2 = manager
        .acquire_gpu(job2, 1)
        .await
        .expect("Should acquire GPU 1");

    // Check updated statistics
    let stats = manager.get_statistics().await;
    assert_eq!(stats.total_gpus, 3);
    assert_eq!(stats.available_gpus, 1); // GPU 2 still available
    assert_eq!(stats.locked_gpus, 2);
    assert_eq!(stats.active_jobs, 2);

    // Verify active jobs list
    let active_jobs = manager
        .list_active_jobs()
        .await
        .expect("Should list active jobs");
    assert_eq!(active_jobs.len(), 2);
    assert!(active_jobs.contains(&(0, job1)));
    assert!(active_jobs.contains(&(1, job2)));
}

#[tokio::test]
async fn test_gpu_manager_tracks_ownership() {
    let manager = GPUResourceManager::new(vec![0, 1])
        .await
        .expect("Failed to create GPU manager");

    let job_id = Uuid::new_v4();

    // Initially, GPU is not locked
    assert!(!manager.is_gpu_locked(0).await);
    assert!(manager.get_gpu_owner(0).await.is_none());

    // Acquire GPU
    let _lock = manager
        .acquire_gpu(job_id, 0)
        .await
        .expect("Should acquire GPU");

    // Verify ownership
    assert!(manager.is_gpu_locked(0).await);
    assert_eq!(manager.get_gpu_owner(0).await, Some(job_id));

    // Other GPU should still be available
    assert!(!manager.is_gpu_locked(1).await);
    assert!(manager.get_gpu_owner(1).await.is_none());
}

// ============================================================================
// Test 8: Checkpoint Retention Policy with Edge Cases
// ============================================================================

#[tokio::test]
async fn test_checkpoint_retention_handles_ties() {
    let pool = setup_test_db().await;
    let test_model_name = "retention_tie_test";

    // Cleanup before test
    let _ = sqlx::query("DELETE FROM ml_model_versions WHERE metadata->>'test_model_name' = $1")
        .bind(test_model_name)
        .execute(&pool)
        .await;

    let retention_policy = RetentionPolicy {
        max_checkpoints_per_model: 3,
        ranking_metric: "sharpe_ratio".to_string(),
        ascending: false,
    };

    let manager = CheckpointManager::new(pool.clone(), retention_policy)
        .await
        .expect("Failed to create checkpoint manager");

    // Create 5 checkpoints with same Sharpe ratio (tie scenario)
    for i in 0..5 {
        let unique_version = format!(
            "1.0.{}",
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_millis()
                + i as u128
        );
        let metadata = ml::checkpoint::CheckpointMetadata {
            checkpoint_id: Uuid::new_v4().to_string(),
            model_type: ml::ModelType::DQN,
            model_name: test_model_name.to_string(),
            version: unique_version,
            created_at: chrono::Utc::now() - chrono::Duration::days(i),
            epoch: Some(10),
            step: Some(1000),
            loss: Some(0.1),
            accuracy: Some(0.9),
            hyperparameters: HashMap::new(),
            metrics: HashMap::from([
                ("sharpe_ratio".to_string(), 1.5), // Same for all!
            ]),
            architecture: HashMap::new(),
            format: ml::checkpoint::CheckpointFormat::Binary,
            compression: ml::checkpoint::CompressionType::LZ4,
            file_size: 1024,
            compressed_size: None,
            checksum: "0".repeat(64),
            tags: vec![],
            custom_metadata: HashMap::new(),
            signature: None,
            signature_algorithm: "none".to_string(),
            signing_key_id: "test".to_string(),
            signed_at: None,
        };

        manager
            .register_checkpoint(metadata)
            .await
            .expect("Failed to register checkpoint");
    }

    // Apply retention policy
    let archived_count = manager
        .apply_retention_policy(ml::ModelType::DQN, test_model_name)
        .await
        .expect("Failed to apply retention policy");

    // Should archive 2 checkpoints (5 - 3 = 2)
    assert_eq!(
        archived_count, 2,
        "Should archive excess checkpoints even with tied metrics"
    );

    // Verify only 3 remain
    let remaining = manager
        .list_checkpoints(ml::ModelType::DQN, test_model_name)
        .await
        .expect("Failed to list checkpoints");

    assert_eq!(
        remaining.len(),
        3,
        "Should have exactly 3 checkpoints remaining"
    );

    // Cleanup
    let _ = sqlx::query("DELETE FROM ml_model_versions WHERE metadata->>'test_model_name' = $1")
        .bind(test_model_name)
        .execute(&pool)
        .await;
}
