//! Embeddings module using Rig framework

use std::sync::Arc;

use serde::{Deserialize, Serialize};
use tokio::sync::Mutex;

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

/// Document chunk to be embedded
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DocumentChunk {
    /// Unique chunk ID
    pub id: String,

    /// Source document path
    pub source_path: String,

    /// Document type (markdown, kcl, nushell, rust)
    pub doc_type: String,

    /// Chunk content (text to embed)
    pub content: String,

    /// Document category for filtering
    pub category: Option<String>,

    /// Metadata (headings, function names, etc.)
    pub metadata: std::collections::HashMap<String, String>,
}

/// Embedded document with vector
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddedDocument {
    /// Chunk ID
    pub id: String,

    /// Source document path
    pub source_path: String,

    /// Document type
    pub doc_type: String,

    /// Original content
    pub content: String,

    /// Embedding vector
    pub embedding: Vec<f32>,

    /// Metadata
    pub metadata: std::collections::HashMap<String, String>,
}

/// OpenAI HTTP client wrapper for embeddings
#[derive(Clone)]
struct OpenAiClient {
    api_key: String,
    model: String,
    dimension: usize,
}

impl OpenAiClient {
    /// Create new OpenAI client
    fn new(api_key: String, model: String, dimension: usize) -> Self {
        Self {
            api_key,
            model,
            dimension,
        }
    }

    /// Generate embedding via OpenAI API
    async fn embed(&self, text: &str) -> Result<Vec<f32>> {
        let client = reqwest::Client::new();

        #[derive(serde::Serialize)]
        struct EmbeddingRequest {
            input: String,
            model: String,
            dimensions: Option<usize>,
        }

        #[derive(serde::Deserialize)]
        struct EmbeddingResponse {
            data: Vec<EmbeddingData>,
        }

        #[derive(serde::Deserialize)]
        struct EmbeddingData {
            embedding: Vec<f32>,
        }

        let request = EmbeddingRequest {
            input: text.to_string(),
            model: self.model.clone(),
            dimensions: if self.dimension == 1536 {
                None
            } else {
                Some(self.dimension)
            },
        };

        let response = client
            .post("https://api.openai.com/v1/embeddings")
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .json(&request)
            .send()
            .await
            .map_err(|e| {
                crate::error::RagError::embedding(format!("Failed to call OpenAI API: {}", e))
            })?;

        if !response.status().is_success() {
            let error_text = response
                .text()
                .await
                .unwrap_or_else(|_| "Unknown error".to_string());
            return Err(crate::error::RagError::embedding(format!(
                "OpenAI API error: {}",
                error_text
            )));
        }

        let embedding_response: EmbeddingResponse = response.json().await.map_err(|e| {
            crate::error::RagError::embedding(format!("Failed to parse OpenAI response: {}", e))
        })?;

        embedding_response
            .data
            .first()
            .map(|d| d.embedding.clone())
            .ok_or_else(|| crate::error::RagError::embedding("No embedding data in response"))
    }
}

/// Embedding engine using OpenAI via HTTP
pub struct EmbeddingEngine {
    config: EmbeddingConfig,
    client: Arc<Mutex<Option<OpenAiClient>>>,
}

impl EmbeddingEngine {
    /// Create a new embedding engine
    pub fn new(config: EmbeddingConfig) -> Result<Self> {
        // Validate configuration
        match config.provider.as_str() {
            "openai" => {
                if config.openai_api_key.is_none() && !config.fallback_local {
                    return Err(crate::error::RagError::config(
                        "OpenAI API key required for OpenAI provider (or enable fallback_local). \
                         Set OPENAI_API_KEY environment variable.",
                    ));
                }
            }
            "ollama" => {
                tracing::info!("Using Ollama for local embeddings (no API costs)");
            }
            "local" => {
                tracing::info!("Using local embedding model (no API costs)");
            }
            _ => {
                return Err(crate::error::RagError::config(format!(
                    "Unknown embedding provider: {}. Supported: openai, ollama, local",
                    config.provider
                )));
            }
        }

        tracing::info!(
            "Initialized embedding engine: {} (provider: {}, dimension: {})",
            config.model,
            config.provider,
            config.dimension
        );

        Ok(Self {
            config,
            client: Arc::new(Mutex::new(None)),
        })
    }

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

    /// Embed a single chunk using Rig framework
    ///
    /// This is a stub implementation that creates zero vectors.
    /// In production, this would call the Rig embeddings API.
    pub async fn embed_chunk(&self, chunk: &DocumentChunk) -> Result<EmbeddedDocument> {
        let embedding = self.generate_embedding(&chunk.content).await?;

        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(),
        })
    }

    /// Embed multiple chunks in batch
    pub async fn embed_batch(&self, chunks: &[DocumentChunk]) -> Result<Vec<EmbeddedDocument>> {
        let mut results = Vec::new();

        for (idx, chunk) in chunks.iter().enumerate() {
            let embedded = self.embed_chunk(chunk).await?;
            results.push(embedded);

            // Log progress
            if (idx + 1) % 10 == 0 {
                tracing::debug!("Embedded {}/{} chunks", idx + 1, chunks.len());
            }
        }

        tracing::info!("Embedded {} chunks total", chunks.len());

        Ok(results)
    }

    /// Generate embedding for text using configured provider
    async fn generate_embedding(&self, text: &str) -> Result<Vec<f32>> {
        if text.is_empty() {
            return Err(crate::error::RagError::embedding(
                "Empty text cannot be embedded",
            ));
        }

        tracing::trace!("Generating embedding for text: {} chars", text.len());

        match self.config.provider.as_str() {
            "openai" => self.embed_openai(text).await,
            "ollama" => self.embed_ollama(text).await,
            "local" => self.embed_local(text).await,
            _ => Err(crate::error::RagError::embedding(format!(
                "Unknown embedding provider: {}",
                self.config.provider
            ))),
        }
    }

    /// Generate embedding via OpenAI API
    async fn embed_openai(&self, text: &str) -> Result<Vec<f32>> {
        let api_key = self
            .config
            .openai_api_key
            .clone()
            .or_else(|| std::env::var("OPENAI_API_KEY").ok())
            .ok_or_else(|| {
                crate::error::RagError::embedding(
                    "OpenAI API key not found. Set OPENAI_API_KEY env var or config.",
                )
            })?;

        let mut client_lock = self.client.lock().await;
        if client_lock.is_none() {
            *client_lock = Some(OpenAiClient::new(
                api_key.clone(),
                self.config.model.clone(),
                self.config.dimension,
            ));
        }

        let client = client_lock.as_ref().unwrap();
        client.embed(text).await
    }

    /// Generate embedding via Ollama (local)
    async fn embed_ollama(&self, text: &str) -> Result<Vec<f32>> {
        let client = reqwest::Client::new();

        #[derive(serde::Serialize)]
        struct OllamaRequest {
            model: String,
            prompt: String,
        }

        #[derive(serde::Deserialize)]
        struct OllamaResponse {
            embedding: Vec<f32>,
        }

        let request = OllamaRequest {
            model: self.config.model.clone(),
            prompt: text.to_string(),
        };

        let response = client
            .post("http://localhost:11434/api/embeddings")
            .json(&request)
            .send()
            .await
            .map_err(|e| {
                crate::error::RagError::embedding(format!(
                    "Failed to call Ollama API (ensure Ollama is running on localhost:11434): {}",
                    e
                ))
            })?;

        if !response.status().is_success() {
            let error_text = response
                .text()
                .await
                .unwrap_or_else(|_| "Unknown error".to_string());
            return Err(crate::error::RagError::embedding(format!(
                "Ollama API error: {}",
                error_text
            )));
        }

        let embedding_response: OllamaResponse = response.json().await.map_err(|e| {
            crate::error::RagError::embedding(format!("Failed to parse Ollama response: {}", e))
        })?;

        Ok(embedding_response.embedding)
    }

    /// Generate embedding using local model (stub for future implementation)
    async fn embed_local(&self, text: &str) -> Result<Vec<f32>> {
        tracing::warn!(
            "Local embeddings not fully implemented. Returning zero vector for now. For \
             production local embeddings, use Ollama or integrate huggingface transformers."
        );

        // Return zero vector of correct dimension
        // Future: integrate sentence-transformers or ONNX models
        Ok(vec![0.0; self.config.dimension])
    }
}

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

    #[test]
    fn test_embedding_engine_creation_openai() {
        let config = EmbeddingConfig {
            provider: "openai".to_string(),
            openai_api_key: Some("test-key".to_string()),
            ..Default::default()
        };

        let result = EmbeddingEngine::new(config);
        assert!(result.is_ok());
    }

    #[test]
    fn test_embedding_engine_creation_openai_no_key() {
        let config = EmbeddingConfig {
            provider: "openai".to_string(),
            openai_api_key: None,
            fallback_local: false,
            ..Default::default()
        };

        let result = EmbeddingEngine::new(config);
        assert!(result.is_err());
    }

    #[test]
    fn test_local_embedding_engine() {
        let config = EmbeddingConfig {
            provider: "local".to_string(),
            ..Default::default()
        };

        let engine = EmbeddingEngine::new(config);
        assert!(engine.is_ok());
    }

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

        let engine = EmbeddingEngine::new(config).unwrap();

        let chunk = DocumentChunk {
            id: "test-1".to_string(),
            source_path: "test.md".to_string(),
            doc_type: "markdown".to_string(),
            content: "This is test content".to_string(),
            category: None,
            metadata: std::collections::HashMap::new(),
        };

        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(), 1536); // Default dimension
    }
}
