//! Additional KV Cache Tests (Part 02) - Coverage Enhancement
//!
//! Targets uncovered paths in gpu/scheduler/kv.rs:
//! - RoPE application edge cases
//! - Error handling paths
//! - Sampling strategies (argmax, top-k)
//! - SwiGLU vs GELU FFN branches
//! - GQA incremental attention paths
//! - Large sequence and edge dimension cases

use crate::gpu::executor::MockExecutor;
use crate::gpu::scheduler::{
    BlockWeights, GpuGenerateConfig, GpuModel, GpuModelConfig,
};
use crate::gpu::StreamingKVCache;

// ============================================================================
// Test Helpers
// ============================================================================

/// Create a minimal config for fast testing
fn minimal_config() -> GpuModelConfig {
    GpuModelConfig {
        hidden_dim: 32,
        intermediate_dim: 64,
        num_layers: 1,
        num_heads: 2,
        num_kv_heads: 2,
        vocab_size: 50,
        eps: 1e-5,
        rope_theta: 10000.0,
    }
}

/// Create a GQA config (num_heads > num_kv_heads)
fn gqa_config() -> GpuModelConfig {
    GpuModelConfig {
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 2,
        num_heads: 8,
        num_kv_heads: 2, // 4 Q heads per KV head
        vocab_size: 100,
        eps: 1e-5,
        rope_theta: 10000.0,
    }
}

/// Create config with SwiGLU FFN (needs gate weights)
fn swiglu_config() -> GpuModelConfig {
    GpuModelConfig {
        hidden_dim: 32,
        intermediate_dim: 64,
        num_layers: 1,
        num_heads: 2,
        num_kv_heads: 2,
        vocab_size: 50,
        eps: 1e-6,
        rope_theta: 10000.0,
    }
}

/// Create model with SwiGLU gate weights
fn create_swiglu_model() -> GpuModel {
    let config = swiglu_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    // Add gate weights to enable SwiGLU path
    if !model.block_weights.is_empty() {
        let gate_weight = vec![0.01f32; config.hidden_dim * config.intermediate_dim];
        model.block_weights[0].ffn_gate_weight = Some(gate_weight);
    }

    model
}

// ============================================================================
// forward_gpu_with_cache Error Path Tests
// ============================================================================

#[test]
fn test_forward_gpu_with_cache_empty_tokens_error() {
    let config = minimal_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    let result = model.forward_gpu_with_cache(&[], &mut kv_cache);
    assert!(result.is_err());

    let err = result.unwrap_err();
    let err_msg = format!("{:?}", err);
    assert!(err_msg.contains("empty") || err_msg.contains("Empty"));
}

#[test]
fn test_forward_gpu_with_cache_out_of_bounds_token() {
    let config = minimal_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    // Token 9999 exceeds vocab_size=50
    let result = model.forward_gpu_with_cache(&[9999], &mut kv_cache);
    assert!(result.is_err());

    let err = result.unwrap_err();
    let err_msg = format!("{:?}", err);
    assert!(err_msg.contains("out of bounds") || err_msg.contains("Token"));
}

// ============================================================================
// forward_gpu_incremental Error Path Tests
// ============================================================================

#[test]
fn test_forward_gpu_incremental_out_of_bounds_token() {
    let config = minimal_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    // Token 9999 exceeds vocab_size=50
    let result = model.forward_gpu_incremental(9999, &mut kv_cache);
    assert!(result.is_err());
}

#[test]
fn test_forward_gpu_incremental_boundary_token() {
    let config = minimal_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("boundary_token");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    // Token at boundary (vocab_size - 1 = 49)
    let result = model.forward_gpu_incremental(49, &mut kv_cache);
    assert!(result.is_ok());
}

// ============================================================================
// generate_with_cache Error Path Tests
// ============================================================================

#[test]
fn test_generate_with_cache_empty_prompt_error() {
    let config = minimal_config();
    let mut model = GpuModel::new(config).unwrap();

    let gen_config = GpuGenerateConfig::deterministic(5);
    let result = model.generate_with_cache(&[], &gen_config);

    assert!(result.is_err());
    let err_msg = format!("{:?}", result.unwrap_err());
    assert!(err_msg.contains("empty") || err_msg.contains("Prompt"));
}

#[test]
fn test_generate_with_cache_single_token_prompt() {
    let config = minimal_config();
    let mut model = GpuModel::new(config).unwrap();

    let mock = MockExecutor::new("single_prompt");
    model.with_test_executor(Box::new(mock));

    let gen_config = GpuGenerateConfig::deterministic(2);
    let result = model.generate_with_cache(&[1], &gen_config);

    assert!(result.is_ok());
    let tokens = result.unwrap();
    assert!(!tokens.is_empty());
}

// ============================================================================
// Sampling Strategy Tests (argmax and sample_topk paths)
// ============================================================================

#[test]
fn test_generate_with_cache_greedy_sampling() {
    let config = minimal_config();
    let mut model = GpuModel::new(config).unwrap();

    let mock = MockExecutor::new("greedy");
    model.with_test_executor(Box::new(mock));

    // temperature=0.0 triggers argmax path
    let gen_config = GpuGenerateConfig::deterministic(3);
    assert_eq!(gen_config.temperature, 0.0);
    assert_eq!(gen_config.top_k, 1);

    let result = model.generate_with_cache(&[1, 2], &gen_config);
    assert!(result.is_ok());
}

#[test]
fn test_generate_with_cache_topk_sampling() {
    let config = minimal_config();
    let mut model = GpuModel::new(config).unwrap();

    let mock = MockExecutor::new("topk");
    model.with_test_executor(Box::new(mock));

    // temperature > 0 and top_k > 1 triggers sample_topk path
    let gen_config = GpuGenerateConfig::with_sampling(3, 0.8, 5);
    assert!(gen_config.temperature > 0.0);
    assert!(gen_config.top_k > 1);

    let result = model.generate_with_cache(&[1], &gen_config);
    assert!(result.is_ok());
}

#[test]
fn test_generate_with_cache_high_temperature() {
    let config = minimal_config();
    let mut model = GpuModel::new(config).unwrap();

    let mock = MockExecutor::new("high_temp");
    model.with_test_executor(Box::new(mock));

    // High temperature increases entropy in sampling
    let gen_config = GpuGenerateConfig::with_sampling(2, 1.5, 10);
    let result = model.generate_with_cache(&[1], &gen_config);

    assert!(result.is_ok());
}

#[test]
fn test_generate_with_cache_topk_1_with_temperature() {
    let config = minimal_config();
    let mut model = GpuModel::new(config).unwrap();

    let mock = MockExecutor::new("topk_1");
    model.with_test_executor(Box::new(mock));

    // top_k=1 should also trigger argmax path
    let gen_config = GpuGenerateConfig {
        max_tokens: 2,
        temperature: 0.7,
        top_k: 1,
        stop_tokens: vec![],
        trace: false,
    };

    let result = model.generate_with_cache(&[1], &gen_config);
    assert!(result.is_ok());
}

// ============================================================================
// Stop Token Tests
// ============================================================================

#[test]
fn test_generate_with_cache_stop_token_first_iteration() {
    let config = minimal_config();
    let mut model = GpuModel::new(config).unwrap();

    // MockExecutor returns zeros, so argmax returns 0
    let mock = MockExecutor::new("stop_first");
    model.with_test_executor(Box::new(mock));

    // Stop on token 0 (which is what argmax of zeros returns)
    let gen_config = GpuGenerateConfig::deterministic(10).with_stop_tokens(vec![0]);
    let result = model.generate_with_cache(&[1], &gen_config);

    assert!(result.is_ok());
    let tokens = result.unwrap();
    // Should stop early due to stop token
    assert!(tokens.len() <= 2); // prompt + at most 1 generated
}

#[test]
fn test_generate_with_cache_multiple_stop_tokens() {
    let config = minimal_config();
    let mut model = GpuModel::new(config).unwrap();

    let mock = MockExecutor::new("multi_stop");
    model.with_test_executor(Box::new(mock));

    let gen_config = GpuGenerateConfig::deterministic(5)
        .with_stop_tokens(vec![0, 1, 2, 3]);

    let result = model.generate_with_cache(&[5], &gen_config);
    assert!(result.is_ok());
}

// ============================================================================
// GQA (Grouped Query Attention) Path Tests
// ============================================================================

#[test]
fn test_forward_gpu_with_cache_gqa() {
    let config = gqa_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("gqa_cache");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    let result = model.forward_gpu_with_cache(&[1, 2, 3], &mut kv_cache);
    assert!(result.is_ok());
}

#[test]
fn test_forward_gpu_incremental_gqa() {
    let config = gqa_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("gqa_incremental");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    // First populate cache
    let _ = model.forward_gpu_with_cache(&[1], &mut kv_cache);

    // Then incremental
    let result = model.forward_gpu_incremental(2, &mut kv_cache);
    assert!(result.is_ok());
}

#[test]
fn test_generate_with_cache_gqa_multiple_iterations() {
    let config = gqa_config();
    let mut model = GpuModel::new(config).unwrap();

    let mock = MockExecutor::new("gqa_gen");
    model.with_test_executor(Box::new(mock));

    let gen_config = GpuGenerateConfig::deterministic(5);
    let result = model.generate_with_cache(&[1, 2], &gen_config);

    assert!(result.is_ok());
    let tokens = result.unwrap();
    assert!(tokens.len() >= 2); // At least the prompt
}

// ============================================================================
// SwiGLU FFN Path Tests
// ============================================================================

#[test]
fn test_forward_gpu_with_cache_swiglu_path() {
    let mut model = create_swiglu_model();
    let config = model.config().clone();

    let mock = MockExecutor::new("swiglu");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    let result = model.forward_gpu_with_cache(&[1, 2], &mut kv_cache);
    assert!(result.is_ok());
}

#[test]
fn test_forward_gpu_incremental_swiglu_path() {
    let mut model = create_swiglu_model();
    let config = model.config().clone();

    let mock = MockExecutor::new("swiglu_inc");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    // Initialize cache
    let _ = model.forward_gpu_with_cache(&[1], &mut kv_cache);

    // Incremental with SwiGLU
    let result = model.forward_gpu_incremental(2, &mut kv_cache);
    assert!(result.is_ok());
}

// ============================================================================
// Multi-token Sequence Tests (RoPE with multiple positions)
// ============================================================================

#[test]
fn test_forward_gpu_with_cache_long_sequence() {
    let config = minimal_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("long_seq");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        128,
        config.num_kv_heads,
        config.head_dim(),
    );

    // Test with longer sequence to exercise RoPE over multiple positions
    let tokens: Vec<usize> = (0..10).collect();
    let result = model.forward_gpu_with_cache(&tokens, &mut kv_cache);

    assert!(result.is_ok());
}

#[test]
fn test_incremental_multiple_steps() {
    let config = minimal_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("multi_step");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    // Initialize
    let _ = model.forward_gpu_with_cache(&[1], &mut kv_cache);

    // Multiple incremental steps to test RoPE position advancement
    for i in 2..8 {
        let result = model.forward_gpu_incremental(i, &mut kv_cache);
        assert!(result.is_ok(), "Failed at step {}", i);
    }
}

// ============================================================================
// KV Cache State Tests
// ============================================================================

#[test]
fn test_kv_cache_population_during_forward() {
    let config = minimal_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("cache_pop");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    // Initially cache should be empty
    let (k, v) = kv_cache.get_valid(0);
    assert!(k.is_empty());
    assert!(v.is_empty());

    // Forward should populate cache
    let _ = model.forward_gpu_with_cache(&[1, 2, 3], &mut kv_cache);

    let (k, v) = kv_cache.get_valid(0);
    // Cache should now have 3 positions worth of K/V
    let kv_dim = config.num_kv_heads * config.head_dim();
    assert_eq!(k.len(), 3 * kv_dim);
    assert_eq!(v.len(), 3 * kv_dim);
}

#[test]
fn test_kv_cache_growth_during_incremental() {
    let config = minimal_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("cache_grow");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    let kv_dim = config.num_kv_heads * config.head_dim();

    // Initialize with 1 token
    let _ = model.forward_gpu_with_cache(&[1], &mut kv_cache);
    let (k1, _) = kv_cache.get_valid(0);
    assert_eq!(k1.len(), kv_dim);

    // Incremental should grow cache
    let _ = model.forward_gpu_incremental(2, &mut kv_cache);
    let (k2, _) = kv_cache.get_valid(0);
    assert_eq!(k2.len(), 2 * kv_dim);
}

// ============================================================================
// Edge Case Tests
// ============================================================================

#[test]
fn test_forward_single_token_only() {
    let config = minimal_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("single");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    let result = model.forward_gpu_with_cache(&[0], &mut kv_cache);
    assert!(result.is_ok());

    let logits = result.unwrap();
    assert_eq!(logits.len(), config.vocab_size);
}

#[test]
fn test_generate_max_tokens_zero() {
    let config = minimal_config();
    let mut model = GpuModel::new(config).unwrap();

    let mock = MockExecutor::new("zero_max");
    model.with_test_executor(Box::new(mock));

    let gen_config = GpuGenerateConfig::deterministic(0);
    let result = model.generate_with_cache(&[1, 2], &gen_config);

    assert!(result.is_ok());
    let tokens = result.unwrap();
    // Should return at least prompt
    assert!(!tokens.is_empty());
}

#[test]
fn test_generate_with_all_stop_tokens() {
    let config = minimal_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("all_stop");
    model.with_test_executor(Box::new(mock));

    // All possible tokens as stop tokens
    let stop_tokens: Vec<usize> = (0..config.vocab_size).collect();
    let gen_config = GpuGenerateConfig::deterministic(10)
        .with_stop_tokens(stop_tokens);

    let result = model.generate_with_cache(&[25], &gen_config);
    assert!(result.is_ok());

    let tokens = result.unwrap();
    // Should stop immediately after first generation
    assert!(tokens.len() <= 2);
}

// ============================================================================
// Multi-layer Tests
// ============================================================================

#[test]
fn test_forward_multi_layer_model() {
    let config = GpuModelConfig {
        hidden_dim: 32,
        intermediate_dim: 64,
        num_layers: 4, // Multiple layers
        num_heads: 2,
        num_kv_heads: 2,
        vocab_size: 50,
        eps: 1e-5,
        rope_theta: 10000.0,
    };

    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("multi_layer");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    let result = model.forward_gpu_with_cache(&[1, 2], &mut kv_cache);
    assert!(result.is_ok());

    // Verify all layers have cached values
    for layer in 0..config.num_layers {
        let (k, v) = kv_cache.get_valid(layer);
        assert!(!k.is_empty(), "Layer {} K cache empty", layer);
        assert!(!v.is_empty(), "Layer {} V cache empty", layer);
    }
}
