//! Integration tests for fastembed embeddings

#[cfg(feature = "embeddings")]
#[tokio::test]
async fn test_fastembed_initialization() {
    use fastembed::TextEmbedding;

    let model = TextEmbedding::try_new(Default::default());

    assert!(model.is_ok(), "Failed to initialize fastembed model: {:?}", model.err());
}

#[cfg(feature = "embeddings")]
#[tokio::test]
async fn test_fastembed_embedding_generation() {
    use fastembed::TextEmbedding;

    let mut model = TextEmbedding::try_new(Default::default()).expect("Failed to initialize model");

    let texts = vec![
        "Hello world, this is a test.",
        "Fastembed is a Rust embedding library.",
        "Testing embedding generation stability.",
    ];

    let result = model.embed(texts.clone(), None);
    assert!(result.is_ok(), "Failed to generate embeddings: {:?}", result.err());

    let embeddings = result.expect("Operation failed");
    assert_eq!(embeddings.len(), 3, "Expected 3 embeddings");

    for (i, embedding) in embeddings.iter().enumerate() {
        assert_eq!(embedding.len(), 384, "Embedding {} has wrong dimensions", i);

        let sum: f32 = embedding.iter().sum();
        assert!(sum.abs() > 0.0001, "Embedding {} appears to be all zeros", i);
    }

    println!(
        "✓ Successfully generated {} embeddings with 384 dimensions each",
        embeddings.len()
    );
}

#[cfg(feature = "embeddings")]
#[tokio::test]
async fn test_fastembed_batch_processing() {
    use fastembed::TextEmbedding;

    let mut model = TextEmbedding::try_new(Default::default()).expect("Failed to initialize model");

    let texts: Vec<String> = (0..50)
        .map(|i| {
            format!(
                "This is test sentence number {}. It contains some text for embedding.",
                i
            )
        })
        .collect();

    let start = std::time::Instant::now();
    let result = model.embed(texts.clone(), Some(32));
    let duration = start.elapsed();

    assert!(result.is_ok(), "Batch embedding failed: {:?}", result.err());

    let embeddings = result.expect("Operation failed");
    assert_eq!(embeddings.len(), 50, "Expected 50 embeddings");

    println!(
        "✓ Batch processed 50 texts in {:?} ({:.2} ms per text)",
        duration,
        duration.as_millis() as f64 / 50.0
    );
}

#[cfg(feature = "embeddings")]
#[tokio::test]
async fn test_fastembed_different_models() {
    use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};

    let models_to_test = vec![
        (EmbeddingModel::AllMiniLML6V2Q, 384, "AllMiniLML6V2Q (fast, quantized)"),
        (EmbeddingModel::BGEBaseENV15, 768, "BGEBaseENV15 (balanced)"),
    ];

    let test_text = vec!["Hello world"];

    for (model_type, expected_dims, description) in models_to_test {
        println!("Testing {}", description);

        let model = TextEmbedding::try_new(InitOptions::new(model_type));

        match model {
            Ok(mut m) => {
                let result = m.embed(test_text.clone(), None);
                assert!(result.is_ok(), "Failed to generate embedding for {}", description);

                let embeddings = result.expect("Operation failed");
                assert_eq!(embeddings.len(), 1);
                assert_eq!(
                    embeddings[0].len(),
                    expected_dims,
                    "Wrong dimensions for {}",
                    description
                );

                println!("  ✓ {} produces {}-dim embeddings", description, expected_dims);
            }
            Err(e) => {
                println!("  ⚠ Failed to initialize {}: {:?}", description, e);
            }
        }
    }
}

#[cfg(feature = "embeddings")]
#[tokio::test]
async fn test_fastembed_error_handling() {
    use fastembed::TextEmbedding;

    let mut model = TextEmbedding::try_new(Default::default()).expect("Failed to initialize model");

    let empty_texts: Vec<String> = vec![];
    let result = model.embed(empty_texts, None);

    match result {
        Ok(embeddings) => assert_eq!(embeddings.len(), 0, "Empty input should produce empty output"),
        Err(_) => {
            println!("  ℹ fastembed returns error for empty input (acceptable)");
        }
    }
}

#[cfg(feature = "embeddings")]
#[tokio::test]
async fn test_generate_embeddings_for_chunks_basic() {
    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
    use kreuzberg::embeddings::generate_embeddings_for_chunks;
    use kreuzberg::types::{Chunk, ChunkMetadata};

    let mut chunks = vec![
        Chunk {
            content: "Hello world, this is the first chunk.".to_string(),
            embedding: None,
            metadata: ChunkMetadata {
                byte_start: 0,
                byte_end: 38,
                chunk_index: 0,
                total_chunks: 1,
                token_count: None,
                first_page: None,
                last_page: None,
            },
        },
        Chunk {
            content: "This is the second chunk with different content.".to_string(),
            embedding: None,
            metadata: ChunkMetadata {
                byte_start: 39,
                byte_end: 87,
                chunk_index: 1,
                total_chunks: 1,
                token_count: None,
                first_page: None,
                last_page: None,
            },
        },
        Chunk {
            content: "And this is the third and final chunk.".to_string(),
            embedding: None,
            metadata: ChunkMetadata {
                byte_start: 88,
                byte_end: 126,
                chunk_index: 2,
                total_chunks: 1,
                token_count: None,
                first_page: None,
                last_page: None,
            },
        },
    ];

    let config = EmbeddingConfig {
        model: EmbeddingModelType::Preset {
            name: "fast".to_string(),
        },
        batch_size: 32,
        normalize: false,
        show_download_progress: false,
        cache_dir: None,
    };

    let result = generate_embeddings_for_chunks(&mut chunks, &config);
    assert!(result.is_ok(), "Failed to generate embeddings: {:?}", result.err());

    for (i, chunk) in chunks.iter().enumerate() {
        assert!(chunk.embedding.is_some(), "Chunk {} missing embedding", i);

        let embedding = chunk.embedding.as_ref().expect("Operation failed");
        assert_eq!(embedding.len(), 384, "Chunk {} has wrong embedding dimensions", i);

        let sum: f32 = embedding.iter().sum();
        assert!(sum.abs() > 0.0001, "Chunk {} embedding appears to be all zeros", i);
    }

    println!("✓ Generated embeddings for {} chunks", chunks.len());
}

#[cfg(feature = "embeddings")]
#[tokio::test]
async fn test_generate_embeddings_for_chunks_normalization() {
    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
    use kreuzberg::embeddings::generate_embeddings_for_chunks;
    use kreuzberg::types::{Chunk, ChunkMetadata};

    let test_text = "This is a test sentence for normalization testing.";

    let mut chunks_no_norm = vec![Chunk {
        content: test_text.to_string(),
        embedding: None,
        metadata: ChunkMetadata {
            byte_start: 0,
            byte_end: test_text.len(),
            chunk_index: 0,
            total_chunks: 1,
            token_count: None,
            first_page: None,
            last_page: None,
        },
    }];

    let config_no_norm = EmbeddingConfig {
        model: EmbeddingModelType::Preset {
            name: "fast".to_string(),
        },
        batch_size: 32,
        normalize: false,
        show_download_progress: false,
        cache_dir: None,
    };

    generate_embeddings_for_chunks(&mut chunks_no_norm, &config_no_norm)
        .expect("Failed to generate non-normalized embeddings");

    let mut chunks_norm = vec![Chunk {
        content: test_text.to_string(),
        embedding: None,
        metadata: ChunkMetadata {
            byte_start: 0,
            byte_end: test_text.len(),
            chunk_index: 0,
            total_chunks: 1,
            token_count: None,
            first_page: None,
            last_page: None,
        },
    }];

    let config_norm = EmbeddingConfig {
        model: EmbeddingModelType::Preset {
            name: "fast".to_string(),
        },
        batch_size: 32,
        normalize: true,
        show_download_progress: false,
        cache_dir: None,
    };

    generate_embeddings_for_chunks(&mut chunks_norm, &config_norm).expect("Failed to generate normalized embeddings");

    let embedding_no_norm = chunks_no_norm[0].embedding.as_ref().expect("Operation failed");
    let embedding_norm = chunks_norm[0].embedding.as_ref().expect("Operation failed");

    let magnitude_no_norm: f32 = embedding_no_norm.iter().map(|x| x * x).sum::<f32>().sqrt();
    let magnitude_norm: f32 = embedding_norm.iter().map(|x| x * x).sum::<f32>().sqrt();

    assert!(
        magnitude_no_norm > 0.9 && magnitude_no_norm < 1.1,
        "Embedding magnitude should be reasonable (got {})",
        magnitude_no_norm
    );

    assert!(
        (magnitude_norm - 1.0).abs() < 0.001,
        "Normalized embedding should have unit magnitude (got {})",
        magnitude_norm
    );

    println!(
        "✓ Normalization works: magnitude before={:.4}, after={:.4}",
        magnitude_no_norm, magnitude_norm
    );
}

#[cfg(feature = "embeddings")]
#[tokio::test]
async fn test_generate_embeddings_for_chunks_empty_input() {
    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
    use kreuzberg::embeddings::generate_embeddings_for_chunks;
    use kreuzberg::types::Chunk;

    let mut empty_chunks: Vec<Chunk> = vec![];

    let config = EmbeddingConfig {
        model: EmbeddingModelType::Preset {
            name: "fast".to_string(),
        },
        batch_size: 32,
        normalize: false,
        show_download_progress: false,
        cache_dir: None,
    };

    let result = generate_embeddings_for_chunks(&mut empty_chunks, &config);
    assert!(result.is_ok(), "Empty input should be handled gracefully");

    println!("✓ Empty input handled correctly");
}

#[cfg(feature = "embeddings")]
#[tokio::test]
async fn test_generate_embeddings_for_chunks_model_caching() {
    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
    use kreuzberg::embeddings::generate_embeddings_for_chunks;
    use kreuzberg::types::{Chunk, ChunkMetadata};

    let mut chunks1 = vec![Chunk {
        content: "First batch of text.".to_string(),
        embedding: None,
        metadata: ChunkMetadata {
            byte_start: 0,
            byte_end: 20,
            chunk_index: 0,
            total_chunks: 1,
            token_count: None,
            first_page: None,
            last_page: None,
        },
    }];

    let config = EmbeddingConfig {
        model: EmbeddingModelType::Preset {
            name: "fast".to_string(),
        },
        batch_size: 32,
        normalize: false,
        show_download_progress: false,
        cache_dir: None,
    };

    let start1 = std::time::Instant::now();
    generate_embeddings_for_chunks(&mut chunks1, &config).expect("First call failed");
    let duration1 = start1.elapsed();

    let mut chunks2 = vec![Chunk {
        content: "Second batch of text.".to_string(),
        embedding: None,
        metadata: ChunkMetadata {
            byte_start: 0,
            byte_end: 21,
            chunk_index: 0,
            total_chunks: 1,
            token_count: None,
            first_page: None,
            last_page: None,
        },
    }];

    let start2 = std::time::Instant::now();
    generate_embeddings_for_chunks(&mut chunks2, &config).expect("Second call failed");
    let duration2 = start2.elapsed();

    println!(
        "✓ Model caching works: first call={:?}, second call={:?}",
        duration1, duration2
    );

    assert!(chunks1[0].embedding.is_some());
    assert!(chunks2[0].embedding.is_some());
}

#[cfg(feature = "embeddings")]
#[tokio::test]
async fn test_generate_embeddings_for_chunks_invalid_preset() {
    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
    use kreuzberg::embeddings::generate_embeddings_for_chunks;
    use kreuzberg::types::{Chunk, ChunkMetadata};

    let mut chunks = vec![Chunk {
        content: "Test content".to_string(),
        embedding: None,
        metadata: ChunkMetadata {
            byte_start: 0,
            byte_end: 12,
            chunk_index: 0,
            total_chunks: 1,
            token_count: None,
            first_page: None,
            last_page: None,
        },
    }];

    let config = EmbeddingConfig {
        model: EmbeddingModelType::Preset {
            name: "nonexistent_preset".to_string(),
        },
        batch_size: 32,
        normalize: false,
        show_download_progress: false,
        cache_dir: None,
    };

    let result = generate_embeddings_for_chunks(&mut chunks, &config);
    assert!(result.is_err(), "Should return error for unknown preset");

    let error = result.unwrap_err();
    let error_msg = format!("{}", error);
    assert!(
        error_msg.contains("Unknown embedding preset"),
        "Error should mention unknown preset, got: {}",
        error_msg
    );

    println!("✓ Invalid preset error handling works");
}

#[cfg(feature = "embeddings")]
#[tokio::test]
async fn test_generate_embeddings_for_chunks_unknown_model() {
    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
    use kreuzberg::embeddings::generate_embeddings_for_chunks;
    use kreuzberg::types::{Chunk, ChunkMetadata};

    let mut chunks = vec![Chunk {
        content: "Test content".to_string(),
        embedding: None,
        metadata: ChunkMetadata {
            byte_start: 0,
            byte_end: 12,
            chunk_index: 0,
            total_chunks: 1,
            token_count: None,
            first_page: None,
            last_page: None,
        },
    }];

    let config = EmbeddingConfig {
        model: EmbeddingModelType::FastEmbed {
            model: "UnknownModelXYZ123".to_string(),
            dimensions: 384,
        },
        batch_size: 32,
        normalize: false,
        show_download_progress: false,
        cache_dir: None,
    };

    let result = generate_embeddings_for_chunks(&mut chunks, &config);
    assert!(result.is_err(), "Should return error for unknown model");

    let error = result.unwrap_err();
    let error_msg = format!("{}", error);
    assert!(
        error_msg.contains("Unknown fastembed model"),
        "Error should mention unknown model, got: {}",
        error_msg
    );

    println!("✓ Unknown model error handling works");
}

#[cfg(feature = "embeddings")]
#[tokio::test]
async fn test_generate_embeddings_for_chunks_custom_model_not_supported() {
    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
    use kreuzberg::embeddings::generate_embeddings_for_chunks;
    use kreuzberg::types::{Chunk, ChunkMetadata};

    let mut chunks = vec![Chunk {
        content: "Test content".to_string(),
        embedding: None,
        metadata: ChunkMetadata {
            byte_start: 0,
            byte_end: 12,
            chunk_index: 0,
            total_chunks: 1,
            token_count: None,
            first_page: None,
            last_page: None,
        },
    }];

    let config = EmbeddingConfig {
        model: EmbeddingModelType::Custom {
            model_id: "hf://custom/model".to_string(),
            dimensions: 768,
        },
        batch_size: 32,
        normalize: false,
        show_download_progress: false,
        cache_dir: None,
    };

    let result = generate_embeddings_for_chunks(&mut chunks, &config);
    assert!(result.is_err(), "Should return error for custom models");

    let error = result.unwrap_err();
    let error_msg = format!("{}", error);
    assert!(
        error_msg.contains("Custom ONNX models are not yet supported"),
        "Error should mention custom models not supported, got: {}",
        error_msg
    );

    println!("✓ Custom model error handling works");
}

#[cfg(feature = "embeddings")]
#[tokio::test]
async fn test_generate_embeddings_for_chunks_batch_size() {
    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
    use kreuzberg::embeddings::generate_embeddings_for_chunks;
    use kreuzberg::types::{Chunk, ChunkMetadata};

    let mut chunks: Vec<Chunk> = (0..10)
        .map(|i| Chunk {
            content: format!("This is test chunk number {}.", i),
            embedding: None,
            metadata: ChunkMetadata {
                byte_start: i * 30,
                byte_end: (i + 1) * 30,
                chunk_index: i,
                total_chunks: 10,
                token_count: None,
                first_page: None,
                last_page: None,
            },
        })
        .collect();

    let config = EmbeddingConfig {
        model: EmbeddingModelType::Preset {
            name: "fast".to_string(),
        },
        batch_size: 10,
        normalize: false,
        show_download_progress: false,
        cache_dir: None,
    };

    let result = generate_embeddings_for_chunks(&mut chunks, &config);
    assert!(result.is_ok(), "Processing failed: {:?}", result.err());

    for (i, chunk) in chunks.iter().enumerate() {
        assert!(
            chunk.embedding.is_some(),
            "Chunk {} missing embedding after batch processing",
            i
        );
        assert_eq!(
            chunk.embedding.as_ref().expect("Operation failed").len(),
            384,
            "Chunk {} has wrong dimensions",
            i
        );
    }

    println!("✓ Processing with batch_size=10 works for 10 chunks");
}

#[cfg(all(feature = "embeddings", feature = "chunking"))]
#[tokio::test]
async fn test_generate_embeddings_chunking_integration() {
    use kreuzberg::chunking::{ChunkingConfig, chunk_text};
    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
    use kreuzberg::embeddings::generate_embeddings_for_chunks;

    let text = "This is a test document. It has multiple sentences. \
                Each sentence should be chunked appropriately. \
                The chunking system should create overlapping chunks. \
                Finally, we will generate embeddings for each chunk.";

    let chunking_config = ChunkingConfig {
        max_characters: 50,
        overlap: 10,
        ..Default::default()
    };

    let mut chunking_result = chunk_text(text, &chunking_config, None).expect("Chunking failed");

    assert!(
        chunking_result.chunks.len() > 1,
        "Should create multiple chunks from text"
    );
    println!("✓ Created {} chunks from text", chunking_result.chunks.len());

    let embedding_config = EmbeddingConfig {
        model: EmbeddingModelType::Preset {
            name: "fast".to_string(),
        },
        batch_size: 32,
        normalize: true,
        show_download_progress: false,
        cache_dir: None,
    };

    let result = generate_embeddings_for_chunks(&mut chunking_result.chunks, &embedding_config);
    assert!(result.is_ok(), "Embedding generation failed: {:?}", result.err());

    for (i, chunk) in chunking_result.chunks.iter().enumerate() {
        assert!(chunk.embedding.is_some(), "Chunk {} missing embedding", i);

        let embedding = chunk.embedding.as_ref().expect("Operation failed");
        assert_eq!(embedding.len(), 384, "Chunk {} has wrong embedding dimensions", i);

        let magnitude: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!(
            (magnitude - 1.0).abs() < 0.001,
            "Chunk {} embedding not normalized (magnitude={})",
            i,
            magnitude
        );
    }

    println!(
        "✓ Chunking + Embedding integration works: {} chunks with normalized embeddings",
        chunking_result.chunks.len()
    );
}
