turboprop 0.1.2

//! Integration tests for model management and embedding generation.
//!
//! These tests validate model loading, caching, and embedding generation
//! across different backends with proper isolation and real interactions.
//!
//! Run these tests with: `cargo test model_integration`

use anyhow::Result;
use std::collections::HashMap;
use tempfile::TempDir;
use turboprop::embeddings::{EmbeddingConfig, EmbeddingOptions};
use turboprop::models::ModelManager;
use turboprop::types::ModelName;

/// Test model manager caching functionality
#[tokio::test]
async fn test_model_manager_caching() -> Result<()> {
    let temp_dir = TempDir::new()?;
    let manager = ModelManager::new_with_defaults(temp_dir.path());

    // Initialize cache directory
    manager.init_cache()?;
    assert!(temp_dir.path().exists());

    // Test fastembed model (should be cached after first use)
    let model_name = "sentence-transformers/all-MiniLM-L6-v2";
    let model_name_obj = ModelName::from(model_name);

    // Initially not cached
    assert!(!manager.is_model_cached(&model_name_obj));

    // Verify model exists in available models
    let models = manager.get_available_models();
    let model_info = models
        .iter()
        .find(|m| m.name == model_name_obj)
        .expect("Default model should be available");

    // Test cache path generation
    let cache_path = manager.get_model_path(&model_name_obj);
    let expected_path = temp_dir
        .path()
        .join("sentence-transformers_all-MiniLM-L6-v2");
    assert_eq!(cache_path, expected_path);

    // Validate model info
    assert!(model_info.validate().is_ok());
    assert_eq!(model_info.dimensions, 384);
    assert!(model_info.size_bytes > 0);

    Ok(())
}

/// Test embedding generation consistency across multiple calls
#[tokio::test]
async fn test_embedding_generation_consistency() -> Result<()> {
    let test_text = vec!["function test() { return 42; }".to_string()];

    // Use mock generator for consistent testing
    let config = EmbeddingConfig::default();
    let mut generator = turboprop::embeddings::MockEmbeddingGenerator::new(config);

    // Generate embeddings multiple times
    let embeddings1 = generator.embed_batch(&test_text)?;
    let embeddings2 = generator.embed_batch(&test_text)?;

    assert_eq!(embeddings1.len(), embeddings2.len());
    assert_eq!(embeddings1[0].len(), embeddings2[0].len());

    // Embeddings should be identical (deterministic)
    assert_eq!(embeddings1[0], embeddings2[0]);

    // Verify dimensions
    assert_eq!(embeddings1[0].len(), 384); // Default embedding dimensions

    Ok(())
}

/// Test model switching between different available models
#[tokio::test]
async fn test_model_switching() -> Result<()> {
    let test_texts = vec![
        "function calculateTotal(items) { return items.reduce((sum, item) => sum + item.price, 0); }".to_string(),
    ];

    let manager = ModelManager::default();
    let models = manager.get_available_models();
    let mut embeddings_by_model = HashMap::new();

    // Test with first two available models (using mock generators for speed)
    for model_info in models.iter().take(2) {
        let config = EmbeddingConfig::with_model(model_info.name.as_str())
            .with_embedding_dimensions(model_info.dimensions);

        let mut generator = turboprop::embeddings::MockEmbeddingGenerator::new(config);

        match generator.embed_batch(&test_texts) {
            Ok(embeddings) => {
                embeddings_by_model.insert(model_info.name.clone(), embeddings);
            }
            Err(e) => {
                // Some models may not be available in test environment
                eprintln!(
                    "Could not generate embeddings for {}: {}",
                    model_info.name, e
                );
            }
        }
    }

    // If we have embeddings from different models, verify they have correct dimensions
    if embeddings_by_model.len() >= 2 {
        let model_names: Vec<_> = embeddings_by_model.keys().collect();
        let emb1 = &embeddings_by_model[model_names[0]][0];
        let emb2 = &embeddings_by_model[model_names[1]][0];

        // For mock generator, we mainly verify that embeddings are generated successfully
        // and have reasonable dimensions. Real implementation would produce different embeddings.
        assert!(
            !emb1.is_empty(),
            "First model should produce non-empty embeddings"
        );
        assert!(
            !emb2.is_empty(),
            "Second model should produce non-empty embeddings"
        );

        // In a real implementation, different models would produce different embeddings:
        // assert_ne!(emb1, emb2, "Different models should produce different embeddings");
    }

    Ok(())
}

/// Test embedding options and configuration
#[tokio::test]
async fn test_embedding_options() -> Result<()> {
    let test_texts = vec![
        "This is a test document for retrieval".to_string(),
        "Another document with different content".to_string(),
    ];

    // Test different embedding options
    let options_variants = vec![
        EmbeddingOptions::default(),
        EmbeddingOptions::with_instruction("Represent this document for retrieval"),
        EmbeddingOptions::without_normalization(),
        EmbeddingOptions::with_max_length(256),
    ];

    let config = EmbeddingConfig::default();
    let mut generator = turboprop::embeddings::MockEmbeddingGenerator::new(config);

    for options in options_variants {
        // Test that generator can handle different options without error
        let result = generator.embed_batch(&test_texts);
        assert!(
            result.is_ok(),
            "Embedding generation should succeed with options: {:?}",
            options
        );

        let embeddings = result.unwrap();
        assert_eq!(embeddings.len(), 2);
        assert_eq!(embeddings[0].len(), 384); // Default dimensions
        assert_eq!(embeddings[1].len(), 384);
    }

    Ok(())
}

/// Test model cache statistics and management
#[test]
fn test_model_cache_stats() -> Result<()> {
    let temp_dir = TempDir::new()?;
    let manager = ModelManager::new_with_defaults(temp_dir.path());

    // Initialize cache
    manager.init_cache()?;

    // Test empty cache stats
    let stats = manager.get_cache_stats()?;
    assert_eq!(stats.model_count, 0);
    assert_eq!(stats.total_size_bytes, 0);
    assert_eq!(stats.format_size(), "0.00 B");

    // Create a fake cached model directory to test stats
    let fake_model_dir = manager.get_model_path(&ModelName::from("fake/model"));
    std::fs::create_dir_all(&fake_model_dir)?;
    std::fs::write(fake_model_dir.join("config.json"), "{}")?;
    std::fs::write(fake_model_dir.join("tokenizer.json"), "{}")?;
    std::fs::write(fake_model_dir.join("model.bin"), vec![0u8; 1024])?; // 1KB file

    // Test stats with fake model
    let stats = manager.get_cache_stats()?;
    assert_eq!(stats.model_count, 1);
    assert!(stats.total_size_bytes > 0);
    assert!(stats.format_size().contains("B") || stats.format_size().contains("KB"));

    // Test cache clearing
    manager.clear_cache()?;
    let stats_after_clear = manager.get_cache_stats()?;
    assert_eq!(stats_after_clear.model_count, 0);

    Ok(())
}

/// Test model validation across different model types
#[test]
fn test_model_validation() -> Result<()> {
    let manager = ModelManager::default();
    let models = manager.get_available_models();

    // Test that all available models pass validation
    for model in &models {
        assert!(
            model.validate().is_ok(),
            "Model {} should pass validation",
            model.name
        );

        // Verify basic properties
        assert!(!model.name.as_str().is_empty());
        assert!(!model.description.is_empty());
        assert!(model.dimensions > 0);
        assert!(model.size_bytes > 0);

        // Type-specific validation
        match model.model_type {
            turboprop::types::ModelType::SentenceTransformer => {
                assert_eq!(model.backend, turboprop::types::ModelBackend::FastEmbed);
            }
            turboprop::types::ModelType::GGUF => {
                assert_eq!(model.backend, turboprop::types::ModelBackend::Candle);
                assert!(
                    model.download_url.is_some(),
                    "GGUF models should have download URL"
                );
            }
            turboprop::types::ModelType::HuggingFace => {
                assert_eq!(model.backend, turboprop::types::ModelBackend::Custom);
            }
        }
    }

    // Verify we have models of each type
    let has_sentence_transformer = models.iter().any(|m| {
        matches!(
            m.model_type,
            turboprop::types::ModelType::SentenceTransformer
        )
    });
    let has_gguf = models
        .iter()
        .any(|m| matches!(m.model_type, turboprop::types::ModelType::GGUF));
    let has_huggingface = models
        .iter()
        .any(|m| matches!(m.model_type, turboprop::types::ModelType::HuggingFace));

    assert!(
        has_sentence_transformer,
        "Should have at least one SentenceTransformer model"
    );
    assert!(has_gguf, "Should have at least one GGUF model");
    assert!(
        has_huggingface,
        "Should have at least one HuggingFace model"
    );

    Ok(())
}

/// Test embedding configuration with different models
#[test]
fn test_embedding_config_models() {
    let manager = ModelManager::default();
    let models = manager.get_available_models();

    for model in models.iter().take(3) {
        // Test first 3 models
        let config = EmbeddingConfig::with_model(model.name.as_str())
            .with_embedding_dimensions(model.dimensions)
            .with_batch_size(16);

        assert_eq!(config.model_name, model.name.as_str());
        assert_eq!(config.embedding_dimensions, model.dimensions);
        assert_eq!(config.batch_size, 16);
        // Cache directory should now default to ~/.turboprop/models
        let expected_cache_dir = dirs::home_dir()
            .map(|p| p.join(".turboprop").join("models"))
            .unwrap_or_else(|| std::path::PathBuf::from(".turboprop/models"));
        assert_eq!(config.cache_dir, expected_cache_dir);
    }
}

/// Test error handling for invalid model names
#[tokio::test]
async fn test_invalid_model_handling() {
    let config = EmbeddingConfig::with_model("nonexistent/model");

    // Mock generator should handle any model name
    let mut generator = turboprop::embeddings::MockEmbeddingGenerator::new(config);

    let test_texts = vec!["test text".to_string()];
    let result = generator.embed_batch(&test_texts);

    // Mock generator should succeed even with invalid model names
    assert!(result.is_ok());
}

/// Test batch processing with different batch sizes
#[tokio::test]
async fn test_batch_processing() -> Result<()> {
    let test_texts: Vec<String> = (0..10)
        .map(|i| format!("Test document number {}", i))
        .collect();

    let batch_sizes = vec![1, 3, 5, 10, 20]; // Test various batch sizes

    for batch_size in batch_sizes {
        let config = EmbeddingConfig::default().with_batch_size(batch_size);
        let mut generator = turboprop::embeddings::MockEmbeddingGenerator::new(config);

        let result = generator.embed_batch(&test_texts);
        assert!(
            result.is_ok(),
            "Batch processing should succeed with batch size {}",
            batch_size
        );

        let embeddings = result.unwrap();
        assert_eq!(embeddings.len(), test_texts.len());

        // Verify all embeddings have correct dimensions
        for (i, embedding) in embeddings.iter().enumerate() {
            assert_eq!(
                embedding.len(),
                384,
                "Embedding {} should have correct dimensions",
                i
            );
        }
    }

    Ok(())
}

/// Test model-specific features
#[test]
fn test_model_specific_features() {
    let manager = ModelManager::default();
    let models = manager.get_available_models();

    // Test GGUF model features
    if let Some(gguf_model) = models
        .iter()
        .find(|m| matches!(m.model_type, turboprop::types::ModelType::GGUF))
    {
        assert!(
            gguf_model.name.as_str().contains("gguf"),
            "GGUF model name should contain 'gguf'"
        );
        assert!(
            gguf_model.description.to_lowercase().contains("code")
                || gguf_model.description.to_lowercase().contains("embedding"),
            "GGUF model should be for code or embeddings"
        );
    }

    // Test Qwen3 model features
    if let Some(qwen_model) = models.iter().find(|m| m.name.as_str().contains("Qwen3")) {
        assert_eq!(
            qwen_model.model_type,
            turboprop::types::ModelType::HuggingFace
        );
        assert_eq!(qwen_model.backend, turboprop::types::ModelBackend::Custom);
        assert!(
            qwen_model.dimensions >= 1024,
            "Qwen3 should have high-dimensional embeddings"
        );
    }

    // Test SentenceTransformer model features
    if let Some(st_model) = models.iter().find(|m| {
        matches!(
            m.model_type,
            turboprop::types::ModelType::SentenceTransformer
        )
    }) {
        assert_eq!(st_model.backend, turboprop::types::ModelBackend::FastEmbed);
        assert!(st_model.name.as_str().contains("sentence-transformers"));
    }
}

/// Helper function to calculate cosine similarity between two vectors
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    assert_eq!(a.len(), b.len(), "Vectors must have the same length");

    let dot_product: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();

    if norm_a == 0.0 || norm_b == 0.0 {
        0.0
    } else {
        dot_product / (norm_a * norm_b)
    }
}

/// Test similarity calculations
#[test]
fn test_cosine_similarity() {
    let vec1 = vec![1.0, 0.0, 0.0];
    let vec2 = vec![1.0, 0.0, 0.0];
    let vec3 = vec![0.0, 1.0, 0.0];

    assert!(
        (cosine_similarity(&vec1, &vec2) - 1.0).abs() < 1e-6,
        "Identical vectors should have similarity 1.0"
    );
    assert!(
        (cosine_similarity(&vec1, &vec3) - 0.0).abs() < 1e-6,
        "Orthogonal vectors should have similarity 0.0"
    );
}

/// Test embedding similarity for similar and different texts
#[tokio::test]
async fn test_embedding_similarity_patterns() -> Result<()> {
    let similar_texts = vec![
        "function calculateSum(a, b) { return a + b; }".to_string(),
        "function addNumbers(x, y) { return x + y; }".to_string(),
    ];

    let different_texts = vec![
        "function calculateSum(a, b) { return a + b; }".to_string(),
        "The weather is nice today".to_string(),
    ];

    let config = EmbeddingConfig::default();
    let mut generator = turboprop::embeddings::MockEmbeddingGenerator::new(config);

    let similar_embeddings = generator.embed_batch(&similar_texts)?;
    let different_embeddings = generator.embed_batch(&different_texts)?;

    let similar_similarity = cosine_similarity(&similar_embeddings[0], &similar_embeddings[1]);
    let different_similarity =
        cosine_similarity(&different_embeddings[0], &different_embeddings[1]);

    // Since we're using mock embeddings, we can't test actual semantic similarity,
    // but we can verify the calculations work and produce valid results
    assert!(
        (-1.0..=1.0).contains(&similar_similarity),
        "Similarity should be in [-1, 1] range"
    );
    assert!(
        (-1.0..=1.0).contains(&different_similarity),
        "Similarity should be in [-1, 1] range"
    );

    Ok(())
}