//! Embeddings module using stratum-embeddings
//!
//! Provides unified embedding providers (OpenAI, Ollama, HuggingFace,
//! FastEmbed) with caching, batch processing, and fallback support.

use std::sync::Arc;

use async_trait::async_trait;
// Import the trait so methods are available
use stratum_embeddings::EmbeddingProvider as StratumProvider;
// Re-export stratum-embeddings types
pub use stratum_embeddings::{
    EmbeddingOptions, HuggingFaceModel, HuggingFaceProvider, OllamaModel, OllamaProvider,
    OpenAiModel, OpenAiProvider,
};
use tracing::debug;

pub type Result<T> = std::result::Result<T, EmbeddingError>;

#[derive(Debug, thiserror::Error)]
pub enum EmbeddingError {
    #[error("Configuration error: {0}")]
    ConfigError(String),

    #[error("Request failed: {0}")]
    RequestFailed(String),

    #[error("Invalid input: {0}")]
    InvalidInput(String),

    #[error("Stratum embeddings error: {0}")]
    StratumError(#[from] stratum_embeddings::EmbeddingError),
}

/// Embedding provider trait (vapora compatibility layer)
#[async_trait]
pub trait EmbeddingProvider: Send + Sync {
    async fn embed(&self, text: &str) -> Result<Vec<f32>>;
    fn provider_name(&self) -> &str;
    fn model_name(&self) -> &str;
}

// ============================================================================
// Provider Wrappers (delegate to stratum-embeddings)
// ============================================================================

pub struct OllamaEmbedding {
    provider: OllamaProvider,
}

impl OllamaEmbedding {
    pub fn new(_endpoint: String, model: String) -> Result<Self> {
        // Note: stratum-embeddings OllamaProvider uses default endpoint
        // The endpoint parameter is kept for API compatibility but not used
        let model_enum = if model == "nomic-embed-text" {
            OllamaModel::NomicEmbed
        } else if model == "mxbai-embed-large" {
            OllamaModel::MxbaiEmbed
        } else if model == "all-minilm" {
            OllamaModel::AllMiniLm
        } else {
            // Custom model - default dimensions to 768
            OllamaModel::Custom(model, 768)
        };

        let provider = OllamaProvider::new(model_enum)?;
        Ok(Self { provider })
    }
}

#[async_trait]
impl EmbeddingProvider for OllamaEmbedding {
    async fn embed(&self, text: &str) -> Result<Vec<f32>> {
        let options = EmbeddingOptions::default_with_cache();
        Ok(self.provider.embed(text, &options).await?)
    }

    fn provider_name(&self) -> &str {
        "ollama"
    }

    fn model_name(&self) -> &str {
        self.provider.model()
    }
}

pub struct OpenAIEmbedding {
    provider: OpenAiProvider,
}

impl OpenAIEmbedding {
    pub fn new(api_key: String, model: String) -> Result<Self> {
        let model_enum = match model.as_str() {
            "text-embedding-3-small" => OpenAiModel::TextEmbedding3Small,
            "text-embedding-3-large" => OpenAiModel::TextEmbedding3Large,
            "text-embedding-ada-002" => OpenAiModel::TextEmbeddingAda002,
            _ => {
                debug!(
                    "Unknown OpenAI model '{}', defaulting to text-embedding-3-small",
                    model
                );
                OpenAiModel::TextEmbedding3Small
            }
        };

        let provider = OpenAiProvider::new(api_key, model_enum)?;
        Ok(Self { provider })
    }
}

#[async_trait]
impl EmbeddingProvider for OpenAIEmbedding {
    async fn embed(&self, text: &str) -> Result<Vec<f32>> {
        let options = EmbeddingOptions::default_with_cache();
        Ok(self.provider.embed(text, &options).await?)
    }

    fn provider_name(&self) -> &str {
        "openai"
    }

    fn model_name(&self) -> &str {
        self.provider.model()
    }
}

pub struct HuggingFaceEmbedding {
    provider: HuggingFaceProvider,
}

impl HuggingFaceEmbedding {
    pub fn new(api_key: String, model: String) -> Result<Self> {
        let model_enum = match model.as_str() {
            "BAAI/bge-small-en-v1.5" => HuggingFaceModel::BgeSmall,
            "BAAI/bge-base-en-v1.5" => HuggingFaceModel::BgeBase,
            "BAAI/bge-large-en-v1.5" => HuggingFaceModel::BgeLarge,
            "sentence-transformers/all-MiniLM-L6-v2" => HuggingFaceModel::AllMiniLm,
            "sentence-transformers/all-mpnet-base-v2" => HuggingFaceModel::AllMpnet,
            _ => {
                debug!("Using custom HuggingFace model: {}", model);
                HuggingFaceModel::Custom(model, 384) // Default to 384 dims
            }
        };

        let provider = HuggingFaceProvider::new(api_key, model_enum)?;
        Ok(Self { provider })
    }
}

#[async_trait]
impl EmbeddingProvider for HuggingFaceEmbedding {
    async fn embed(&self, text: &str) -> Result<Vec<f32>> {
        let options = EmbeddingOptions::default_with_cache();
        Ok(self.provider.embed(text, &options).await?)
    }

    fn provider_name(&self) -> &str {
        "huggingface"
    }

    fn model_name(&self) -> &str {
        self.provider.model()
    }
}

// ============================================================================
// Factory function to create providers from environment/config
// ============================================================================

pub async fn create_embedding_provider(provider_name: &str) -> Result<Arc<dyn EmbeddingProvider>> {
    match provider_name.to_lowercase().as_str() {
        "ollama" => {
            let endpoint = std::env::var("OLLAMA_ENDPOINT")
                .unwrap_or_else(|_| "http://localhost:11434".to_string());
            let model = std::env::var("OLLAMA_EMBEDDING_MODEL")
                .unwrap_or_else(|_| "nomic-embed-text".to_string());

            debug!("Creating Ollama embedding provider: {}", model);
            Ok(Arc::new(OllamaEmbedding::new(endpoint, model)?))
        }

        "openai" => {
            let api_key = std::env::var("OPENAI_API_KEY")
                .map_err(|_| EmbeddingError::ConfigError("OPENAI_API_KEY not set".to_string()))?;
            let model = std::env::var("OPENAI_EMBEDDING_MODEL")
                .unwrap_or_else(|_| "text-embedding-3-small".to_string());

            debug!("Creating OpenAI embedding provider: {}", model);
            Ok(Arc::new(OpenAIEmbedding::new(api_key, model)?))
        }

        "huggingface" => {
            let api_key = std::env::var("HUGGINGFACE_API_KEY")
                .or_else(|_| std::env::var("HF_TOKEN"))
                .map_err(|_| {
                    EmbeddingError::ConfigError(
                        "HUGGINGFACE_API_KEY or HF_TOKEN not set".to_string(),
                    )
                })?;
            let model = std::env::var("HUGGINGFACE_MODEL")
                .unwrap_or_else(|_| "BAAI/bge-small-en-v1.5".to_string());

            debug!("Creating HuggingFace embedding provider: {}", model);
            Ok(Arc::new(HuggingFaceEmbedding::new(api_key, model)?))
        }

        _ => Err(EmbeddingError::ConfigError(format!(
            "Unknown embedding provider: {}. Supported: ollama, openai, huggingface",
            provider_name
        ))),
    }
}

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

    #[test]
    fn test_ollama_provider_creation() {
        let ollama = OllamaEmbedding::new(
            "http://localhost:11434".to_string(),
            "nomic-embed-text".to_string(),
        );
        assert!(ollama.is_ok());
        let ollama = ollama.unwrap();
        assert_eq!(ollama.provider_name(), "ollama");
        assert_eq!(ollama.model_name(), "nomic-embed-text");
    }

    #[test]
    fn test_openai_provider_creation() {
        let openai =
            OpenAIEmbedding::new("test-key".to_string(), "text-embedding-3-small".to_string());
        assert!(openai.is_ok());
        let openai = openai.unwrap();
        assert_eq!(openai.provider_name(), "openai");
        assert_eq!(openai.model_name(), "text-embedding-3-small");
    }

    #[test]
    fn test_huggingface_provider_creation() {
        let hf =
            HuggingFaceEmbedding::new("test-key".to_string(), "BAAI/bge-small-en-v1.5".to_string());
        assert!(hf.is_ok());
        let hf = hf.unwrap();
        assert_eq!(hf.provider_name(), "huggingface");
        assert_eq!(hf.model_name(), "BAAI/bge-small-en-v1.5");
    }

    #[test]
    fn test_empty_text_error() {
        let embedding_error = EmbeddingError::InvalidInput("Empty text".to_string());
        assert!(embedding_error.to_string().contains("Empty text"));
    }
}
