//! Tests for Metal Flash Attention implementation
//!
//! Following strict TDD - tests verify Metal backend behavior

#![cfg(all(test, feature = "metal", target_os = "macos"))]

use super::metal::{FlashMetalKernels, MetalKernelResult};
use crate::config::FlashAttentionConfig;
use crate::error::FlashResult;
use objc2_metal::MTLCreateSystemDefaultDevice;

/// Test: Metal device availability on macOS
#[test]
fn test_metal_device_available() {
    let device = unsafe { MTLCreateSystemDefaultDevice() };
    assert!(device.is_some(), "Metal device must be available on macOS");
}

/// Test: FlashMetalKernels creation succeeds
#[test]
fn test_metal_kernels_creation() {
    let config = FlashAttentionConfig::new(8, 64);
    let result = FlashMetalKernels::new(&config);

    assert!(result.is_ok(), "Metal kernels creation should succeed");

    let kernels = result.unwrap();
    assert_eq!(kernels.config().num_heads, 8);
    assert_eq!(kernels.config().head_dim, 64);
}

/// Test: Device info retrieval
#[test]
fn test_get_device_info() {
    let config = FlashAttentionConfig::new(8, 64);
    let kernels = FlashMetalKernels::new(&config).unwrap();

    let device_info = kernels.get_device_info();
    assert!(device_info.is_ok());

    let info_str = device_info.unwrap();
    assert!(info_str.contains("Metal Device"));
}

/// Test: Block size optimization
#[test]
fn test_optimize_block_sizes() {
    let config = FlashAttentionConfig::new(8, 64);
    let kernels = FlashMetalKernels::new(&config).unwrap();

    let test_cases = vec![
        (512, 64, 128, 128),   // seq_len=512 -> should be clamped to 128
        (1024, 128, 128, 128), // seq_len=1024 -> should be 128
        (256, 64, 128, 128),   // seq_len=256 -> should be 128
        (64, 64, 64, 64),      // seq_len=64 -> should be 64
    ];

    for (seq_len, head_dim, expected_q, expected_kv) in test_cases {
        let (block_q, block_kv) = kernels.optimize_block_sizes(seq_len, head_dim);

        // Validate block sizes are within reasonable bounds
        assert!(block_q >= 32, "Block Q must be at least 32");
        assert!(block_q <= 256, "Block Q must not exceed 256");
        assert!(block_kv >= 32, "Block KV must be at least 32");
        assert!(block_kv <= 256, "Block KV must not exceed 256");

        // Check alignment to warp size (32 for Metal)
        assert_eq!(block_q % 32, 0, "Block Q should be aligned to 32");
        assert_eq!(block_kv % 32, 0, "Block KV should be aligned to 32");
    }
}

/// Test: Multiple kernel instances can coexist
#[test]
fn test_multiple_kernel_instances() {
    let config1 = FlashAttentionConfig::new(8, 64);
    let config2 = FlashAttentionConfig::new(16, 128);

    let kernels1 = FlashMetalKernels::new(&config1);
    let kernels2 = FlashMetalKernels::new(&config2);

    assert!(kernels1.is_ok());
    assert!(kernels2.is_ok());

    let k1 = kernels1.unwrap();
    let k2 = kernels2.unwrap();

    assert_eq!(k1.config().num_heads, 8);
    assert_eq!(k2.config().num_heads, 16);
    assert_eq!(k1.config().head_dim, 64);
    assert_eq!(k2.config().head_dim, 128);
}

/// Test: Forward pass stub returns expected structure
#[tokio::test]
async fn test_forward_pass_stub_structure() {
    use rtx_tensor::{Tensor, Device};

    let config = FlashAttentionConfig::new(8, 64);
    let kernels = FlashMetalKernels::new(&config).unwrap();

    let device = Device::Metal(0); // Assuming Metal device 0
    let batch_size = 2;
    let num_heads = 8;
    let seq_len = 128;
    let head_dim = 64;

    // Create test tensors (would be on Metal device in production)
    let q = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
    let k = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
    let v = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
    let mut output = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
    let mut lse = Tensor::zeros(&[batch_size, num_heads, seq_len], &device).unwrap();

    let result = kernels.flash_attention_forward(
        &q, &k, &v, &mut output, &mut lse, 0.125, false
    ).await;

    // Forward pass should succeed with real execution metrics
    assert!(result.is_ok(), "Forward pass should succeed");

    let kernel_result = result.unwrap();
    // With fixed threadgroup memory (32x32 blocks = 28KB < 32KB limit), kernel should execute
    assert!(kernel_result.execution_time_us > 0, "Kernel should have executed (time > 0)");
    assert!(kernel_result.memory_throughput >= 0.0, "Memory throughput should be non-negative");
}

/// Test: Backward pass execution (with 16x16 blocks = 20KB < 32KB limit)
#[tokio::test]
async fn test_backward_pass_execution() {
    use rtx_tensor::{Tensor, Device};

    let config = FlashAttentionConfig::new(8, 64);
    let kernels = FlashMetalKernels::new(&config).unwrap();

    let device = Device::Metal(0);
    let batch_size = 2;
    let num_heads = 8;
    let seq_len = 128;
    let head_dim = 64;

    let dout = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
    let q = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
    let k = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
    let v = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
    let output = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
    let lse = Tensor::zeros(&[batch_size, num_heads, seq_len], &device).unwrap();
    let mut grad_q = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
    let mut grad_k = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
    let mut grad_v = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();

    let result = kernels.flash_attention_backward(
        &dout, &q, &k, &v, &output, &lse,
        &mut grad_q, &mut grad_k, &mut grad_v,
        0.125, false
    ).await;

    // Backward pass should succeed with 16x16 blocks (20KB < 32KB limit)
    assert!(result.is_ok(), "Backward pass should succeed with fixed block sizes");

    let kernel_result = result.unwrap();
    assert!(kernel_result.execution_time_us > 0, "Backward kernel should have executed");
}

/// Test: Configuration validation
#[test]
fn test_config_validation() {
    let valid_configs = vec![
        FlashAttentionConfig::new(8, 64),
        FlashAttentionConfig::new(16, 128),
        FlashAttentionConfig::new(32, 256),
    ];

    for config in valid_configs {
        let result = FlashMetalKernels::new(&config);
        assert!(result.is_ok(), "Valid config should create kernels successfully");
    }
}

/// Test: Thread safety - kernels can be shared across threads
#[test]
fn test_thread_safety() {
    use std::sync::Arc;
    use std::thread;

    let config = FlashAttentionConfig::new(8, 64);
    let kernels = Arc::new(FlashMetalKernels::new(&config).unwrap());

    let mut handles = vec![];

    for _ in 0..4 {
        let kernels_clone = Arc::clone(&kernels);
        let handle = thread::spawn(move || {
            let (block_q, block_kv) = kernels_clone.optimize_block_sizes(1024, 64);
            assert!(block_q >= 32 && block_q <= 256);
            assert!(block_kv >= 32 && block_kv <= 256);
        });
        handles.push(handle);
    }

    for handle in handles {
        handle.join().unwrap();
    }
}