//! TFT Trainer Tests
//!
//! Unit tests for the Temporal Fusion Transformer trainer.

use super::*;
use crate::checkpoint::FileSystemStorage;
use crate::MLError;
use ml_core::device::MlDevice;
use std::path::PathBuf;
use std::sync::Arc;
use tracing::info;

fn cuda_device() -> MlDevice {
    MlDevice::new_cuda(0).expect("CUDA device required")
}

#[tokio::test]
async fn test_tft_trainer_creation() {
    let config = TFTTrainerConfig::default();
    let storage = Arc::new(FileSystemStorage::new(PathBuf::from(
        "/tmp/test_checkpoints",
    )));

    let trainer = TFTTrainer::new(config, storage);
    assert!(trainer.is_ok());
}

#[tokio::test]
async fn test_training_config_conversion() {
    let config = TFTTrainerConfig {
        hidden_dim: 128,
        num_attention_heads: 4,
        lstm_layers: 2,
        ..Default::default()
    };

    let model_config = config.to_model_config();
    assert_eq!(model_config.hidden_dim, 128);
    assert_eq!(model_config.num_heads, 4);
    assert_eq!(model_config.num_layers, 2);
}

#[tokio::test]
async fn test_checkpoint_save_load() {
    use tempfile::TempDir;

    // Create temporary directory for checkpoints
    let temp_dir = TempDir::new().expect("Failed to create temp dir");
    let checkpoint_dir = temp_dir.path().to_str().unwrap().to_string();

    // Create trainer with custom checkpoint directory
    let config = TFTTrainerConfig {
        epochs: 5,
        batch_size: 2,
        hidden_dim: 32,
        num_attention_heads: 2,
        checkpoint_dir: checkpoint_dir.clone(),
        ..Default::default()
    };

    let storage = Arc::new(FileSystemStorage::new(PathBuf::from(&checkpoint_dir)));
    let trainer = TFTTrainer::new(config, storage).expect("Failed to create trainer");

    // Save checkpoint
    let result = trainer.save_checkpoint(1, 0.5, 0.6).await;
    assert!(
        result.is_ok(),
        "Failed to save checkpoint: {:?}",
        result.err()
    );

    // Verify checkpoint file exists and has non-zero size
    let checkpoint_path = PathBuf::from(&checkpoint_dir).join("tft_225_epoch_1.safetensors");
    assert!(checkpoint_path.exists(), "Checkpoint file does not exist");

    let file_size = std::fs::metadata(&checkpoint_path)
        .expect("Failed to get file metadata")
        .len();
    assert!(
        file_size > 0,
        "Checkpoint file is empty (size: {} bytes)",
        file_size
    );

    // Note: File size will be small (16-32 bytes) for untrained model with empty VarMap
    // In actual training, weights would be present and file size would be >1MB
    // Here we just verify the SafeTensors format is being saved correctly

    // Verify metadata file exists
    let metadata_path = checkpoint_path.with_extension("json");
    assert!(metadata_path.exists(), "Metadata file does not exist");

    // Read and validate metadata
    let metadata_content =
        std::fs::read_to_string(&metadata_path).expect("Failed to read metadata");
    let metadata: serde_json::Value =
        serde_json::from_str(&metadata_content).expect("Failed to parse metadata JSON");

    assert_eq!(metadata["epoch"], 1);
    assert_eq!(metadata["model_type"], "TFT");
    assert!(metadata["metrics"]["train_loss"].as_f64().unwrap() - 0.5 < 0.0001);
    assert!(metadata["metrics"]["val_loss"].as_f64().unwrap() - 0.6 < 0.0001);

    info!(bytes = file_size, "Checkpoint saved successfully");
    info!(path = %metadata_path.display(), "Metadata file created");
}

#[tokio::test]
async fn test_zero_batch_size_handling() {
    use tempfile::TempDir;

    // Create temporary directory
    let temp_dir = TempDir::new().expect("Failed to create temp dir");
    let checkpoint_dir = temp_dir.path().to_str().unwrap().to_string();

    // Test TFT rejects zero batch size
    let config = TFTTrainerConfig {
        batch_size: 0,
        checkpoint_dir,
        ..Default::default()
    };

    let storage = Arc::new(FileSystemStorage::new(PathBuf::from(
        "/tmp/test_checkpoints",
    )));
    let result = TFTTrainer::new(config, storage);

    // Should fail with descriptive error
    assert!(
        result.is_err(),
        "TFT should reject zero batch size, but got: {:?}",
        result
    );

    // Error message should mention batch size or validation
    let error_msg = result.unwrap_err().to_string();
    assert!(
        error_msg.to_lowercase().contains("batch")
            || error_msg.to_lowercase().contains("valid"),
        "Error message should mention batch size or validation, got: {}",
        error_msg
    );
}

#[tokio::test]
async fn test_is_oom_error() {
    // Test standard OOM error strings
    let oom1 = MLError::ModelError("CUDA error: out of memory".to_owned());
    assert!(
        TFTTrainer::is_oom_error(&oom1),
        "Should detect 'out of memory'"
    );

    let oom2 = MLError::TrainingError("OOM detected during forward pass".to_owned());
    assert!(TFTTrainer::is_oom_error(&oom2), "Should detect 'OOM'");

    let oom3 = MLError::ModelError("cuda error 2: allocation failed".to_owned());
    assert!(
        TFTTrainer::is_oom_error(&oom3),
        "Should detect 'cuda error 2'"
    );

    let oom4 = MLError::ModelError("Failed to allocate 500MB on GPU".to_owned());
    assert!(
        TFTTrainer::is_oom_error(&oom4),
        "Should detect 'failed to allocate'"
    );

    // Test non-OOM errors
    let not_oom1 = MLError::ModelError("Invalid tensor shape".to_owned());
    assert!(
        !TFTTrainer::is_oom_error(&not_oom1),
        "Should not detect regular errors"
    );

    let not_oom2 = MLError::ConfigError("Missing parameter".to_owned());
    assert!(
        !TFTTrainer::is_oom_error(&not_oom2),
        "Should not detect config errors"
    );
}

#[tokio::test]
async fn test_sync_cuda_device_cpu() {
    // Test CUDA sync on CPU device (should be no-op)
    let device = cuda_device();
    let result = TFTTrainer::sync_cuda_device(&device);
    assert!(result.is_ok(), "CPU sync should succeed as no-op");
}

#[tokio::test]
async fn test_sync_cuda_device_gpu() {
    // Test CUDA sync on GPU device
    let device = MlDevice::cuda_if_available(0);
    if !device.is_cuda() {
        info!("Skipping CUDA sync test - GPU not available");
        return;
    }

    let result = TFTTrainer::sync_cuda_device(&device);
    assert!(
        result.is_ok(),
        "GPU sync should succeed: {:?}",
        result.err()
    );
}

#[tokio::test]
async fn test_oom_retry_batch_size_reduction() {
    use tempfile::TempDir;

    // Create temporary directory for checkpoints
    let temp_dir = TempDir::new().expect("Failed to create temp dir");
    let checkpoint_dir = temp_dir.path().to_str().unwrap().to_string();

    // Create trainer with initial batch size
    let config = TFTTrainerConfig {
        epochs: 5,
        batch_size: 64, // Start with large batch size
        hidden_dim: 32,
        checkpoint_dir: checkpoint_dir.clone(),
        ..Default::default()
    };

    let storage = Arc::new(FileSystemStorage::new(PathBuf::from(&checkpoint_dir)));
    let mut trainer = TFTTrainer::new(config, storage).expect("Failed to create trainer");

    // Simulate OOM retry logic
    let mut current_batch_size = trainer.training_config.batch_size;
    let mut oom_retry_count = 0;
    const MAX_OOM_RETRIES: usize = 3;

    // Simulate 3 OOM events
    while oom_retry_count < MAX_OOM_RETRIES {
        oom_retry_count += 1;
        current_batch_size /= 2;

        // Test exponential backoff: 64 → 32 → 16 → 8
        match oom_retry_count {
            1 => assert_eq!(
                current_batch_size, 32,
                "First retry should halve batch size to 32"
            ),
            2 => assert_eq!(
                current_batch_size, 16,
                "Second retry should halve batch size to 16"
            ),
            3 => assert_eq!(
                current_batch_size, 8,
                "Third retry should halve batch size to 8"
            ),
            _ => panic!("Should not exceed MAX_OOM_RETRIES"),
        }

        // Update trainer config (simulating actual retry logic)
        trainer.training_config.batch_size = current_batch_size;
        trainer.training_config.validation_batch_size = current_batch_size;
    }

    // Verify final state
    assert_eq!(oom_retry_count, 3, "Should have retried exactly 3 times");
    assert_eq!(current_batch_size, 8, "Final batch size should be 8");
    assert_eq!(
        trainer.training_config.batch_size, 8,
        "Trainer config should be updated"
    );
}

#[tokio::test]
async fn test_oom_retry_minimum_batch_size() {
    use tempfile::TempDir;

    let temp_dir = TempDir::new().expect("Failed to create temp dir");
    let checkpoint_dir = temp_dir.path().to_str().unwrap().to_string();

    // Start with batch size that will go below minimum
    let config = TFTTrainerConfig {
        epochs: 5,
        batch_size: 4, // Minimum batch size
        hidden_dim: 32,
        checkpoint_dir: checkpoint_dir.clone(),
        ..Default::default()
    };

    let storage = Arc::new(FileSystemStorage::new(PathBuf::from(&checkpoint_dir)));
    let trainer = TFTTrainer::new(config, storage).expect("Failed to create trainer");

    // Simulate OOM at minimum batch size
    let mut current_batch_size = trainer.training_config.batch_size;
    current_batch_size /= 2; // 4 → 2

    // Should be below minimum (4)
    assert!(
        current_batch_size < 4,
        "Reduced batch size should be below minimum (got {})",
        current_batch_size
    );
}
