// Embedding Generation - Multi-Provider Support
// Integrates with vapora-llm-router for embedding generation
// Supports: Claude, OpenAI, Gemini, Ollama

use tracing::{debug, info};

/// Embedding provider selection
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EmbeddingProvider {
    /// OpenAI text-embedding-3-small/large
    OpenAI,
    /// Voyage AI (via LLMRouter)
    Voyage,
    /// Cohere embed models
    Cohere,
    /// Local Ollama embeddings
    Ollama,
}

impl std::fmt::Display for EmbeddingProvider {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            EmbeddingProvider::OpenAI => write!(f, "openai"),
            EmbeddingProvider::Voyage => write!(f, "voyage"),
            EmbeddingProvider::Cohere => write!(f, "cohere"),
            EmbeddingProvider::Ollama => write!(f, "ollama"),
        }
    }
}

/// Embedding configuration
#[derive(Debug, Clone)]
pub struct EmbeddingConfig {
    /// Provider to use
    pub provider: EmbeddingProvider,
    /// Model name (provider-specific)
    pub model: String,
    /// Embedding dimensions (provider-specific)
    pub dimensions: usize,
    /// Batch size for embedding requests
    pub batch_size: usize,
}

impl Default for EmbeddingConfig {
    fn default() -> Self {
        Self {
            provider: EmbeddingProvider::OpenAI,
            model: "text-embedding-3-small".to_string(),
            dimensions: 1536,
            batch_size: 100,
        }
    }
}

impl EmbeddingConfig {
    /// Create config for OpenAI text-embedding-3-small
    pub fn openai_small() -> Self {
        Self {
            provider: EmbeddingProvider::OpenAI,
            model: "text-embedding-3-small".to_string(),
            dimensions: 1536,
            batch_size: 100,
        }
    }

    /// Create config for OpenAI text-embedding-3-large
    pub fn openai_large() -> Self {
        Self {
            provider: EmbeddingProvider::OpenAI,
            model: "text-embedding-3-large".to_string(),
            dimensions: 3072,
            batch_size: 100,
        }
    }

    /// Create config for Ollama (local)
    pub fn ollama(model: impl Into<String>) -> Self {
        Self {
            provider: EmbeddingProvider::Ollama,
            model: model.into(),
            dimensions: 768, // Default for most Ollama models
            batch_size: 50,
        }
    }

    /// Create config for Voyage AI
    pub fn voyage(model: impl Into<String>) -> Self {
        Self {
            provider: EmbeddingProvider::Voyage,
            model: model.into(),
            dimensions: 1024, // voyage-2 default
            batch_size: 100,
        }
    }
}

/// Embedding generator - integrates with LLMRouter
pub struct EmbeddingGenerator {
    config: EmbeddingConfig,
    // Phase 5: Simplified implementation - no actual LLMRouter integration yet
    // Real integration will use Arc<LLMRouter> and call embedding endpoints
}

impl EmbeddingGenerator {
    /// Create a new embedding generator
    pub fn new(config: EmbeddingConfig) -> Self {
        Self { config }
    }

    /// Generate embedding for a single text
    ///
    /// # Arguments
    /// - `text`: Text to embed
    ///
    /// # Returns
    /// Embedding vector (dimensions match config)
    pub async fn embed_single(&self, text: &str) -> crate::Result<Vec<f32>> {
        debug!(
            "Generating embedding: provider={}, model={}, text_len={}",
            self.config.provider,
            self.config.model,
            text.len()
        );

        // Phase 5: Simplified implementation - generate deterministic embeddings
        // Real implementation would call LLMRouter embedding endpoints
        let embedding = self.generate_mock_embedding(text);

        Ok(embedding)
    }

    /// Generate embeddings for multiple texts (batched)
    ///
    /// # Arguments
    /// - `texts`: Texts to embed
    ///
    /// # Returns
    /// Vector of embedding vectors
    pub async fn embed_batch(&self, texts: &[String]) -> crate::Result<Vec<Vec<f32>>> {
        info!(
            "Generating embeddings: provider={}, model={}, batch_size={}",
            self.config.provider,
            self.config.model,
            texts.len()
        );

        // Process in batches
        let mut all_embeddings = Vec::new();

        for chunk in texts.chunks(self.config.batch_size) {
            for text in chunk {
                let embedding = self.embed_single(text).await?;
                all_embeddings.push(embedding);
            }
        }

        Ok(all_embeddings)
    }

    /// Generate mock embedding (Phase 5 placeholder)
    ///
    /// Real implementation will call:
    /// - OpenAI: POST https://api.openai.com/v1/embeddings
    /// - Voyage: POST https://api.voyageai.com/v1/embeddings
    /// - Ollama: POST http://localhost:11434/api/embeddings
    fn generate_mock_embedding(&self, text: &str) -> Vec<f32> {
        // Generate deterministic embedding based on text hash
        let mut embedding = vec![0.0; self.config.dimensions];

        // Simple hash-based generation for testing
        let hash = text.chars().enumerate().fold(0u32, |acc, (i, c)| {
            acc.wrapping_add(c as u32 * (i as u32 + 1))
        });

        for (i, val) in embedding.iter_mut().enumerate() {
            let seed = hash.wrapping_add(i as u32);
            *val = (seed as f32 / u32::MAX as f32) * 2.0 - 1.0; // Range: [-1,
                                                                // 1]
        }

        // Normalize to unit length (cosine similarity assumes normalized vectors)
        let magnitude = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
        if magnitude > 0.0 {
            for val in &mut embedding {
                *val /= magnitude;
            }
        }

        embedding
    }

    /// Get embedding configuration
    pub fn config(&self) -> &EmbeddingConfig {
        &self.config
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_embedding_generator_creation() {
        let config = EmbeddingConfig::default();
        let generator = EmbeddingGenerator::new(config);
        assert_eq!(generator.config().provider, EmbeddingProvider::OpenAI);
    }

    #[tokio::test]
    async fn test_embed_single() {
        let config = EmbeddingConfig::openai_small();
        let generator = EmbeddingGenerator::new(config);

        let embedding = generator.embed_single("Hello, world!").await.unwrap();
        assert_eq!(embedding.len(), 1536); // OpenAI small dimensions
    }

    #[tokio::test]
    async fn test_embed_batch() {
        let config = EmbeddingConfig::openai_small();
        let generator = EmbeddingGenerator::new(config);

        let texts = vec![
            "First chunk".to_string(),
            "Second chunk".to_string(),
            "Third chunk".to_string(),
        ];

        let embeddings = generator.embed_batch(&texts).await.unwrap();
        assert_eq!(embeddings.len(), 3);
        assert_eq!(embeddings[0].len(), 1536);
    }

    #[tokio::test]
    async fn test_embedding_normalized() {
        let config = EmbeddingConfig::openai_small();
        let generator = EmbeddingGenerator::new(config);

        let embedding = generator.embed_single("Test text").await.unwrap();

        // Check normalization (unit length)
        let magnitude: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!(
            (magnitude - 1.0).abs() < 0.001,
            "Embedding should be normalized to unit length"
        );
    }

    #[tokio::test]
    async fn test_embedding_deterministic() {
        let config = EmbeddingConfig::openai_small();
        let generator = EmbeddingGenerator::new(config);

        let embedding1 = generator.embed_single("Test text").await.unwrap();
        let embedding2 = generator.embed_single("Test text").await.unwrap();

        assert_eq!(
            embedding1, embedding2,
            "Same text should produce same embedding"
        );
    }

    #[tokio::test]
    async fn test_embedding_different_texts() {
        let config = EmbeddingConfig::openai_small();
        let generator = EmbeddingGenerator::new(config);

        let embedding1 = generator.embed_single("First text").await.unwrap();
        let embedding2 = generator.embed_single("Second text").await.unwrap();

        assert_ne!(
            embedding1, embedding2,
            "Different texts should produce different embeddings"
        );
    }

    #[tokio::test]
    async fn test_config_openai_small() {
        let config = EmbeddingConfig::openai_small();
        assert_eq!(config.provider, EmbeddingProvider::OpenAI);
        assert_eq!(config.model, "text-embedding-3-small");
        assert_eq!(config.dimensions, 1536);
    }

    #[tokio::test]
    async fn test_config_openai_large() {
        let config = EmbeddingConfig::openai_large();
        assert_eq!(config.provider, EmbeddingProvider::OpenAI);
        assert_eq!(config.model, "text-embedding-3-large");
        assert_eq!(config.dimensions, 3072);
    }

    #[tokio::test]
    async fn test_config_ollama() {
        let config = EmbeddingConfig::ollama("llama2");
        assert_eq!(config.provider, EmbeddingProvider::Ollama);
        assert_eq!(config.model, "llama2");
        assert_eq!(config.dimensions, 768);
    }

    #[tokio::test]
    async fn test_config_voyage() {
        let config = EmbeddingConfig::voyage("voyage-2");
        assert_eq!(config.provider, EmbeddingProvider::Voyage);
        assert_eq!(config.model, "voyage-2");
        assert_eq!(config.dimensions, 1024);
    }

    #[tokio::test]
    async fn test_embed_batch_respects_batch_size() {
        let mut config = EmbeddingConfig::openai_small();
        config.batch_size = 2; // Small batch size for testing
        let generator = EmbeddingGenerator::new(config);

        let texts = vec![
            "chunk1".to_string(),
            "chunk2".to_string(),
            "chunk3".to_string(),
            "chunk4".to_string(),
            "chunk5".to_string(),
        ];

        let embeddings = generator.embed_batch(&texts).await.unwrap();
        assert_eq!(embeddings.len(), 5);
    }

    #[test]
    fn test_provider_display() {
        assert_eq!(EmbeddingProvider::OpenAI.to_string(), "openai");
        assert_eq!(EmbeddingProvider::Voyage.to_string(), "voyage");
        assert_eq!(EmbeddingProvider::Cohere.to_string(), "cohere");
        assert_eq!(EmbeddingProvider::Ollama.to_string(), "ollama");
    }
}
