//! Embeddings module using stratum-embeddings

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use serde::{Deserialize, Serialize};
use stratum_embeddings::{
    EmbeddingOptions, EmbeddingService, FastEmbedProvider, MemoryCache, OllamaModel,
    OllamaProvider, OpenAiModel, OpenAiProvider,
};

use crate::config::EmbeddingConfig;
use crate::error::Result;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DocumentChunk {
    pub id: String,
    pub source_path: String,
    pub doc_type: String,
    pub content: String,
    pub category: Option<String>,
    pub metadata: HashMap<String, String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddedDocument {
    pub id: String,
    pub source_path: String,
    pub doc_type: String,
    pub content: String,
    pub embedding: Vec<f32>,
    pub metadata: HashMap<String, String>,
}

enum Service {
    OpenAi(EmbeddingService<OpenAiProvider, MemoryCache>),
    Ollama(EmbeddingService<OllamaProvider, MemoryCache>),
    FastEmbed(EmbeddingService<FastEmbedProvider, MemoryCache>),
}

pub struct EmbeddingEngine {
    config: EmbeddingConfig,
    service: Service,
}

impl EmbeddingEngine {
    pub fn new(config: EmbeddingConfig) -> Result<Self> {
        let cache = MemoryCache::new(1000, Duration::from_secs(3600));

        let service = match config.provider.as_str() {
            "openai" => {
                let api_key = config
                    .openai_api_key
                    .clone()
                    .or_else(|| std::env::var("OPENAI_API_KEY").ok());

                let api_key_str = if let Some(key) = api_key {
                    key
                } else if config.fallback_local {
                    "dummy".to_string() // Will fail, but fallback will take
                                        // over
                } else {
                    return Err(crate::error::RagError::config(
                        "OpenAI API key required. Set OPENAI_API_KEY or enable fallback_local",
                    ));
                };

                let model = match config.model.as_str() {
                    "text-embedding-3-small" => OpenAiModel::TextEmbedding3Small,
                    "text-embedding-3-large" => OpenAiModel::TextEmbedding3Large,
                    "text-embedding-ada-002" => OpenAiModel::TextEmbeddingAda002,
                    _ => {
                        tracing::warn!(
                            "Unknown model '{}', using text-embedding-3-small",
                            config.model
                        );
                        OpenAiModel::TextEmbedding3Small
                    }
                };

                let provider = OpenAiProvider::new(api_key_str, model).map_err(|e| {
                    crate::error::RagError::embedding(format!("OpenAI provider failed: {}", e))
                })?;

                let mut svc = EmbeddingService::new(provider).with_cache(cache);

                if config.fallback_local {
                    tracing::info!("Fallback to local embeddings enabled");
                    let fallback = Arc::new(FastEmbedProvider::small().map_err(|e| {
                        crate::error::RagError::embedding(format!("Fallback failed: {}", e))
                    })?);
                    svc = svc.with_fallback(fallback);
                }

                Service::OpenAi(svc)
            }
            "ollama" => {
                use stratum_embeddings::OllamaModel;

                let model = if config.model == "nomic-embed-text" {
                    OllamaModel::NomicEmbed
                } else if config.model == "mxbai-embed-large" {
                    OllamaModel::MxbaiEmbed
                } else if config.model == "all-minilm" {
                    OllamaModel::AllMiniLm
                } else {
                    OllamaModel::Custom(config.model.clone(), config.dimension)
                };

                let provider = OllamaProvider::new(model).map_err(|e| {
                    crate::error::RagError::embedding(format!("Ollama provider failed: {}", e))
                })?;

                Service::Ollama(EmbeddingService::new(provider).with_cache(cache))
            }
            "local" => {
                tracing::info!("Using FastEmbed local embeddings (no API costs)");
                let provider = FastEmbedProvider::small().map_err(|e| {
                    crate::error::RagError::embedding(format!("FastEmbed provider failed: {}", e))
                })?;

                Service::FastEmbed(EmbeddingService::new(provider).with_cache(cache))
            }
            _ => {
                return Err(crate::error::RagError::config(format!(
                    "Unknown provider: {}. Supported: openai, ollama, local",
                    config.provider
                )));
            }
        };

        tracing::info!(
            "Initialized stratum-embeddings: {} (provider: {}, dim: {})",
            config.model,
            config.provider,
            config.dimension
        );

        Ok(Self { config, service })
    }

    pub fn config(&self) -> &EmbeddingConfig {
        &self.config
    }

    pub async fn embed_chunk(&self, chunk: &DocumentChunk) -> Result<EmbeddedDocument> {
        let options = EmbeddingOptions::default_with_cache();

        let embedding = match &self.service {
            Service::OpenAi(svc) => svc.embed(&chunk.content, &options).await,
            Service::Ollama(svc) => svc.embed(&chunk.content, &options).await,
            Service::FastEmbed(svc) => svc.embed(&chunk.content, &options).await,
        }
        .map_err(|e| crate::error::RagError::embedding(format!("Embedding failed: {}", e)))?;

        Ok(EmbeddedDocument {
            id: chunk.id.clone(),
            source_path: chunk.source_path.clone(),
            doc_type: chunk.doc_type.clone(),
            content: chunk.content.clone(),
            embedding,
            metadata: chunk.metadata.clone(),
        })
    }

    pub async fn embed_batch(&self, chunks: &[DocumentChunk]) -> Result<Vec<EmbeddedDocument>> {
        let options = EmbeddingOptions::default_with_cache();
        let texts: Vec<String> = chunks.iter().map(|c| c.content.clone()).collect();

        let result = match &self.service {
            Service::OpenAi(svc) => svc.embed_batch(texts, &options).await,
            Service::Ollama(svc) => svc.embed_batch(texts, &options).await,
            Service::FastEmbed(svc) => svc.embed_batch(texts, &options).await,
        }
        .map_err(|e| crate::error::RagError::embedding(format!("Batch embed failed: {}", e)))?;

        let results = chunks
            .iter()
            .zip(result.embeddings)
            .map(|(chunk, embedding)| EmbeddedDocument {
                id: chunk.id.clone(),
                source_path: chunk.source_path.clone(),
                doc_type: chunk.doc_type.clone(),
                content: chunk.content.clone(),
                embedding,
                metadata: chunk.metadata.clone(),
            })
            .collect();

        tracing::info!("Embedded {} chunks (stratum-embeddings)", chunks.len());
        Ok(results)
    }
}

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

    #[test]
    fn test_engine_openai() {
        let config = EmbeddingConfig {
            provider: "openai".to_string(),
            openai_api_key: Some("test-key".to_string()),
            ..Default::default()
        };
        assert!(EmbeddingEngine::new(config).is_ok());
    }

    #[test]
    fn test_engine_ollama() {
        let config = EmbeddingConfig {
            provider: "ollama".to_string(),
            model: "nomic-embed-text".to_string(),
            ..Default::default()
        };
        assert!(EmbeddingEngine::new(config).is_ok());
    }

    #[test]
    fn test_engine_local() {
        let config = EmbeddingConfig {
            provider: "local".to_string(),
            dimension: 384,
            ..Default::default()
        };
        assert!(EmbeddingEngine::new(config).is_ok());
    }

    #[tokio::test]
    async fn test_embed_chunk() {
        let config = EmbeddingConfig {
            provider: "local".to_string(),
            dimension: 384,
            ..Default::default()
        };
        let engine = EmbeddingEngine::new(config).unwrap();

        let mut metadata = HashMap::new();
        metadata.insert("section".to_string(), "test".to_string());

        let chunk = DocumentChunk {
            id: "test-1".to_string(),
            source_path: "/test/doc.md".to_string(),
            doc_type: "markdown".to_string(),
            content: "Test document".to_string(),
            category: Some("test".to_string()),
            metadata,
        };

        let result = engine.embed_chunk(&chunk).await;
        assert!(result.is_ok());

        let embedded = result.unwrap();
        assert_eq!(embedded.id, "test-1");
        assert_eq!(embedded.embedding.len(), 384);
    }
}
