//! Phase 33: CUDA Driver Integration Tests for GGUF
//!
//! These tests exercise the CUDA inference path to illuminate:
//! - `gguf/cuda/mod.rs` - OwnedQuantizedModelCuda
//! - `gguf/cuda/forward.rs` - forward_cuda
//! - `gguf/cuda/generation.rs` - generate_cuda
//! - `gguf/cuda/backend.rs` - CudaBackend
//!
//! Run with: cargo test --test driver_cuda_gguf --features cuda -- --test-threads=1

#![cfg(feature = "cuda")]

use realizar::gguf::{
    GGUFConfig, OwnedQKVWeights, OwnedQuantizedLayer, OwnedQuantizedModel, OwnedQuantizedTensor,
    QuantizedGenerateConfig, GGUF_TYPE_Q4_K,
};

// =============================================================================
// Test Model Construction Helpers
// =============================================================================

/// Create Q4_K test data for given dimensions
fn create_q4k_test_data(in_dim: usize, out_dim: usize) -> OwnedQuantizedTensor {
    let super_blocks_per_row = in_dim.div_ceil(256);
    let bytes_per_row = super_blocks_per_row * 144;
    let data_size = out_dim * bytes_per_row;
    let mut data = vec![0u8; data_size];

    for row in 0..out_dim {
        for sb in 0..super_blocks_per_row {
            let offset = row * bytes_per_row + sb * 144;
            // d=1.0 in f16 format
            data[offset..offset + 2].copy_from_slice(&0x3C00_u16.to_le_bytes());
            // dmin=0
            data[offset + 2..offset + 4].copy_from_slice(&0x0000_u16.to_le_bytes());
            // Fill with deterministic pattern
            for i in 4..144 {
                data[offset + i] = ((row + sb + i) % 16) as u8;
            }
        }
    }

    OwnedQuantizedTensor {
        data,
        in_dim,
        out_dim,
        qtype: GGUF_TYPE_Q4_K,
    }
}

/// Create a test model with LLaMA-style architecture
fn create_test_model(config: &GGUFConfig) -> OwnedQuantizedModel {
    let vocab_size = config.vocab_size;
    let hidden_dim = config.hidden_dim;
    let intermediate_dim = config.intermediate_dim;
    let kv_dim = config.num_kv_heads * (hidden_dim / config.num_heads);

    let qkv_out_dim = hidden_dim + 2 * kv_dim;
    let qkv_weight = create_q4k_test_data(hidden_dim, qkv_out_dim);
    let attn_output_weight = create_q4k_test_data(hidden_dim, hidden_dim);

    let ffn_up_weight = create_q4k_test_data(hidden_dim, intermediate_dim);
    let ffn_gate_weight = create_q4k_test_data(hidden_dim, intermediate_dim);
    let ffn_down_weight = create_q4k_test_data(intermediate_dim, hidden_dim);

    let attn_norm_weight = vec![1.0f32; hidden_dim];
    let ffn_norm_weight = vec![1.0f32; hidden_dim];

    let mut layers = Vec::with_capacity(config.num_layers);
    for _ in 0..config.num_layers {
        layers.push(OwnedQuantizedLayer {
            attn_norm_weight: attn_norm_weight.clone(),
            attn_norm_bias: None,
            qkv_weight: OwnedQKVWeights::Fused(qkv_weight.clone()),
            qkv_bias: None,
            attn_output_weight: attn_output_weight.clone(),
            attn_output_bias: None,
            ffn_up_weight: ffn_up_weight.clone(),
            ffn_up_bias: None,
            ffn_down_weight: ffn_down_weight.clone(),
            ffn_down_bias: None,
            ffn_gate_weight: Some(ffn_gate_weight.clone()),
            ffn_gate_bias: None,
            ffn_norm_weight: Some(ffn_norm_weight.clone()),
            ffn_norm_bias: None,
        });
    }

    let token_embedding = vec![0.1f32; vocab_size * hidden_dim];
    let output_norm_weight = vec![1.0f32; hidden_dim];
    let lm_head_weight = create_q4k_test_data(hidden_dim, vocab_size);

    OwnedQuantizedModel::new_for_test(
        config.clone(),
        token_embedding,
        layers,
        output_norm_weight,
        None, // output_norm_bias
        lm_head_weight,
        None, // lm_head_bias
    )
}

// =============================================================================
// CUDA Backend Tests (gguf/cuda/backend.rs)
// =============================================================================

#[test]
fn test_driver_cuda_backend_exists() {
    // Illuminates: CudaBackend struct visibility
    use realizar::gguf::CudaBackend;

    let backend = CudaBackend::new(4, 4, 4, 64);
    assert_eq!(backend.m, 4);
    assert_eq!(backend.n, 4);
    assert_eq!(backend.k, 4);
    assert_eq!(backend.head_dim, 64);
}

// =============================================================================
// OwnedQuantizedModelCuda Tests (gguf/cuda/mod.rs)
// =============================================================================

#[test]
fn test_driver_cuda_model_creation() {
    // Illuminates: OwnedQuantizedModelCuda::new
    use realizar::gguf::OwnedQuantizedModelCuda;

    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 2,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model(&config);
    let result = OwnedQuantizedModelCuda::new(model, 0);

    // CUDA init may fail in some environments - that's expected
    // The test illuminates the code path regardless
    match result {
        Ok(cuda_model) => {
            // Access config through model method
            assert_eq!(cuda_model.model().config.hidden_dim, 64);
        },
        Err(e) => {
            // Log but don't fail - CUDA may not be initialized
            eprintln!("CUDA model creation failed (may be expected): {:?}", e);
        },
    }
}

#[test]
fn test_driver_cuda_model_with_max_seq_len() {
    // Illuminates: OwnedQuantizedModelCuda::with_max_seq_len
    use realizar::gguf::OwnedQuantizedModelCuda;

    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 2,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 1024,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model(&config);
    let result = OwnedQuantizedModelCuda::with_max_seq_len(model, 0, 512);

    match result {
        Ok(cuda_model) => {
            // Successfully created with custom seq len
            assert_eq!(cuda_model.model().config.hidden_dim, 64);
        },
        Err(e) => {
            eprintln!("CUDA model with_max_seq_len failed: {:?}", e);
        },
    }
}

// =============================================================================
// CUDA Forward Pass Tests (gguf/cuda/forward.rs)
// =============================================================================

#[test]
fn test_driver_cuda_forward_basic() {
    // Illuminates: forward_cuda
    use realizar::gguf::OwnedQuantizedModelCuda;

    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model(&config);

    match OwnedQuantizedModelCuda::new(model, 0) {
        Ok(mut cuda_model) => {
            let result = cuda_model.forward_cuda(&[42]);
            match result {
                Ok(logits) => {
                    assert_eq!(logits.len(), config.vocab_size);
                    assert!(logits.iter().all(|x| x.is_finite()));
                },
                Err(e) => {
                    eprintln!("forward_cuda failed: {:?}", e);
                },
            }
        },
        Err(e) => {
            eprintln!("CUDA model creation failed: {:?}", e);
        },
    }
}

#[test]
fn test_driver_cuda_forward_multi_token() {
    // Illuminates: forward_cuda with sequence
    use realizar::gguf::OwnedQuantizedModelCuda;

    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model(&config);

    if let Ok(mut cuda_model) = OwnedQuantizedModelCuda::new(model, 0) {
        let tokens = [1u32, 2, 3, 4, 5];
        if let Ok(logits) = cuda_model.forward_cuda(&tokens) {
            assert_eq!(logits.len(), config.vocab_size);
        }
    }
}

// =============================================================================
// CUDA Generation Tests (gguf/cuda/generation.rs)
// =============================================================================

#[test]
fn test_driver_cuda_generate() {
    // Illuminates: generate_cuda
    use realizar::gguf::OwnedQuantizedModelCuda;

    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model(&config);

    if let Ok(mut cuda_model) = OwnedQuantizedModelCuda::new(model, 0) {
        let gen_config = QuantizedGenerateConfig {
            max_tokens: 5,
            temperature: 1.0,
            top_k: 50,
            stop_tokens: vec![],
            trace: false,
        };

        let prompt = vec![1u32, 2, 3];
        if let Ok(output) = cuda_model.generate_cuda(&prompt, &gen_config) {
            assert!(!output.is_empty());
            assert!(output.len() <= prompt.len() + gen_config.max_tokens);
        }
    }
}

#[test]
fn test_driver_cuda_generate_greedy() {
    // Illuminates: temperature=0 greedy path
    use realizar::gguf::OwnedQuantizedModelCuda;

    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model(&config);

    if let Ok(mut cuda_model) = OwnedQuantizedModelCuda::new(model, 0) {
        let gen_config = QuantizedGenerateConfig {
            max_tokens: 3,
            temperature: 0.0, // Greedy
            top_k: 1,
            stop_tokens: vec![],
            trace: false,
        };

        let prompt = vec![42u32];
        if let Ok(output) = cuda_model.generate_cuda(&prompt, &gen_config) {
            assert!(!output.is_empty());
        }
    }
}

#[test]
fn test_driver_cuda_generate_with_stop_token() {
    // Illuminates: stop token handling
    use realizar::gguf::OwnedQuantizedModelCuda;

    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model(&config);

    if let Ok(mut cuda_model) = OwnedQuantizedModelCuda::new(model, 0) {
        let gen_config = QuantizedGenerateConfig {
            max_tokens: 10,
            temperature: 1.0,
            top_k: 50,
            stop_tokens: vec![50], // Stop on token 50
            trace: false,
        };

        let prompt = vec![1u32];
        let _ = cuda_model.generate_cuda(&prompt, &gen_config);
        // Just exercise the code path
    }
}

// =============================================================================
// GQA Tests with CUDA
// =============================================================================

#[test]
fn test_driver_cuda_gqa_forward() {
    // Illuminates: GQA attention path in CUDA
    use realizar::gguf::OwnedQuantizedModelCuda;

    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 8,
        num_kv_heads: 2, // GQA: 4 heads per KV head
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model(&config);

    if let Ok(mut cuda_model) = OwnedQuantizedModelCuda::new(model, 0) {
        if let Ok(logits) = cuda_model.forward_cuda(&[1, 2, 3]) {
            assert_eq!(logits.len(), config.vocab_size);
        }
    }
}

// =============================================================================
// Config Builder Tests
// =============================================================================

#[test]
fn test_driver_config_builder() {
    // Illuminates: QuantizedGenerateConfig builder methods
    let config = QuantizedGenerateConfig::default()
        .with_max_tokens(100)
        .with_temperature(0.7)
        .with_top_k(40)
        .with_stop_tokens(vec![1, 2, 3]);

    assert_eq!(config.max_tokens, 100);
    assert!((config.temperature - 0.7).abs() < 0.01);
    assert_eq!(config.top_k, 40);
    assert_eq!(config.stop_tokens, vec![1, 2, 3]);
}

// =============================================================================
// Multiple Layer Tests
// =============================================================================

#[test]
fn test_driver_cuda_multi_layer() {
    // Illuminates: Multi-layer CUDA forward
    use realizar::gguf::OwnedQuantizedModelCuda;

    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 4, // Multiple layers
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model(&config);

    if let Ok(mut cuda_model) = OwnedQuantizedModelCuda::new(model, 0) {
        if let Ok(logits) = cuda_model.forward_cuda(&[1, 2]) {
            assert_eq!(logits.len(), config.vocab_size);
            assert!(logits.iter().all(|x| x.is_finite()));
        }
    }
}
