//! RAG (Retrieval-Augmented Generation) orchestration
//!
//! Combines semantic search (embeddings + vector store) with
//! full-text search to retrieve relevant documents for LLM queries.

use crate::error::{ErrorWrapper, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

#[cfg(feature = "ai_backend")]
use super::embeddings::{EmbeddingModel, EmbeddingsService};
#[cfg(feature = "ai_backend")]
use super::indexer::{Document, FullTextIndexer};
#[cfg(feature = "ai_backend")]
use super::vector_store::{VectorData, VectorStore};

/// RAG retrieval result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RetrievalResult {
    /// Document ID
    pub doc_id: String,
    /// Document content
    pub content: String,
    /// Semantic relevance score (0-1)
    pub semantic_score: f32,
    /// Keyword relevance score (0-1)
    pub keyword_score: f32,
    /// Combined relevance score
    pub combined_score: f32,
    /// Source type (semantic or keyword)
    pub source: RetrievalSource,
}

/// Source of retrieval
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RetrievalSource {
    /// Retrieved via semantic search
    Semantic,
    /// Retrieved via keyword search
    Keyword,
    /// Retrieved via both methods
    Hybrid,
}

/// RAG configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RagConfig {
    /// Weight for semantic search results (0.0-1.0)
    pub semantic_weight: f32,
    /// Weight for keyword search results (0.0-1.0)
    pub keyword_weight: f32,
    /// Maximum number of results to return
    pub max_results: usize,
    /// Minimum combined score threshold
    pub min_score: f32,
}

impl Default for RagConfig {
    fn default() -> Self {
        RagConfig {
            semantic_weight: 0.6,
            keyword_weight: 0.4,
            max_results: 5,
            min_score: 0.0,
        }
    }
}

/// RAG system combining semantic and keyword search
#[cfg(feature = "ai_backend")]
pub struct RagSystem {
    embeddings: EmbeddingsService,
    vector_store: VectorStore,
    indexer: FullTextIndexer,
    config: RagConfig,
}

#[cfg(feature = "ai_backend")]
impl RagSystem {
    /// Create a new RAG system
    pub fn new(config: RagConfig) -> Result<Self> {
        let embeddings = EmbeddingsService::new(EmbeddingModel::BgeSmallEn)?;
        let vector_store = VectorStore::new(config.max_results);
        let indexer = FullTextIndexer::new();

        Ok(RagSystem {
            embeddings,
            vector_store,
            indexer,
            config,
        })
    }

    /// Create with custom embedding model
    pub fn with_model(model: EmbeddingModel, config: RagConfig) -> Result<Self> {
        let embeddings = EmbeddingsService::new(model)?;
        let vector_store = VectorStore::new(config.max_results);
        let indexer = FullTextIndexer::new();

        Ok(RagSystem {
            embeddings,
            vector_store,
            indexer,
            config,
        })
    }

    /// Add a document to the RAG system
    pub fn add_document(&mut self, id: String, content: String) -> Result<()> {
        // Index for full-text search
        let doc = Document::new(id.clone(), content.clone());
        self.indexer.index(doc)?;

        // Embed and add to vector store
        let embedding = self.embeddings.embed(&content)?;
        let vector_data = VectorData::new(id, embedding.vector);
        self.vector_store.insert(vector_data)?;

        Ok(())
    }

    /// Add multiple documents
    pub fn add_documents(&mut self, docs: Vec<(String, String)>) -> Result<()> {
        for (id, content) in docs {
            self.add_document(id, content)?;
        }
        Ok(())
    }

    /// Add multiple documents in a single batch operation
    ///
    /// More efficient than add_documents() as it:
    /// - Embeds all documents in a single batch call
    /// - Indexes all documents in the full-text indexer
    /// - Invalidates vector store cache only once
    ///
    /// This is significantly faster for large document sets (100+).
    ///
    /// # Arguments
    ///
    /// * `docs` - Vec of (id, content) tuples to add
    ///
    /// # Returns
    ///
    /// Returns error if any document fails validation or embedding
    pub fn add_documents_batch(&mut self, docs: Vec<(String, String)>) -> Result<()> {
        if docs.is_empty() {
            return Ok(());
        }

        // Index all documents in full-text indexer
        let mut document_map = HashMap::new();
        for (id, content) in &docs {
            let doc = Document::new(id.clone(), content.clone());
            self.indexer.index(doc)?;
            document_map.insert(id.clone(), content.clone());
        }

        // Embed all documents in a single batch call
        let contents: Vec<&str> = docs.iter().map(|(_, content)| content.as_str()).collect();
        let embeddings = self.embeddings.embed_batch(&contents)?;

        // Prepare vector data for batch insertion
        let vector_data: Vec<VectorData> = docs
            .iter()
            .zip(embeddings.iter())
            .map(|((id, _), embedding)| VectorData::new(id.clone(), embedding.vector.clone()))
            .collect();

        // Insert all vectors in a single batch operation
        self.vector_store.insert_batch(vector_data)?;

        Ok(())
    }

    /// Retrieve relevant documents for a query
    pub fn retrieve(&mut self, query: &str) -> Result<Vec<RetrievalResult>> {
        if query.is_empty() {
            return Err(ErrorWrapper::validation_failed("Query cannot be empty"));
        }

        // Semantic search
        let embedding = self.embeddings.embed(query)?;
        let semantic_results = self
            .vector_store
            .search(&embedding.vector, self.config.max_results)?;

        // Keyword search
        let keyword_results = self.indexer.search(query, self.config.max_results)?;

        // Combine results
        let mut combined: HashMap<String, RetrievalResult> = HashMap::new();

        // Add semantic results
        for result in semantic_results.iter() {
            let normalized_score = 1.0 / (1.0 + result.distance); // Convert distance to similarity
            let combined_score = normalized_score * self.config.semantic_weight;

            combined.insert(
                result.id.clone(),
                RetrievalResult {
                    doc_id: result.id.clone(),
                    content: String::new(), // Will be filled from indexer
                    semantic_score: normalized_score,
                    keyword_score: 0.0,
                    combined_score,
                    source: RetrievalSource::Semantic,
                },
            );
        }

        // Add keyword results
        for result in keyword_results.iter() {
            let keyword_score = result.score.min(1.0); // Normalize to 0-1
            let keyword_weighted = keyword_score * self.config.keyword_weight;

            combined
                .entry(result.doc_id.clone())
                .and_modify(|r| {
                    r.keyword_score = keyword_score;
                    r.combined_score = (r.semantic_score * self.config.semantic_weight
                        + keyword_score * self.config.keyword_weight)
                        / (self.config.semantic_weight + self.config.keyword_weight);
                    r.source = RetrievalSource::Hybrid;
                    r.content = result.content.clone();
                })
                .or_insert_with(|| RetrievalResult {
                    doc_id: result.doc_id.clone(),
                    content: result.content.clone(),
                    semantic_score: 0.0,
                    keyword_score,
                    combined_score: keyword_weighted,
                    source: RetrievalSource::Keyword,
                });
        }

        // Filter and sort results
        let mut results: Vec<_> = combined
            .into_values()
            .filter(|r| r.combined_score >= self.config.min_score)
            .collect();

        results.sort_by(|a, b| {
            b.combined_score
                .partial_cmp(&a.combined_score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Limit results
        results.truncate(self.config.max_results);

        Ok(results)
    }

    /// Remove a document from the RAG system
    pub fn remove_document(&mut self, id: &str) -> bool {
        let indexer_removed = self.indexer.remove(id).is_some();
        let store_removed = self.vector_store.remove(id).is_some();
        indexer_removed || store_removed
    }

    /// Remove multiple documents in a single batch operation
    ///
    /// More efficient than calling remove_document() multiple times as it:
    /// - Removes all documents from the full-text indexer
    /// - Removes all vectors from the vector store in a single batch
    /// - Invalidates vector store cache only once
    ///
    /// # Arguments
    ///
    /// * `ids` - Slice of document IDs to remove
    ///
    /// # Returns
    ///
    /// Returns the number of documents actually removed
    pub fn remove_documents_batch(&mut self, ids: &[&str]) -> usize {
        let mut count = 0;

        // Remove from full-text indexer
        for id in ids {
            if self.indexer.remove(id).is_some() {
                count += 1;
            }
        }

        // Remove from vector store in batch (single cache invalidation)
        let _store_count = self.vector_store.remove_batch(ids);

        count
    }

    /// Get document count
    pub fn doc_count(&self) -> usize {
        self.indexer.doc_count()
    }

    /// Get RAG configuration
    pub fn config(&self) -> &RagConfig {
        &self.config
    }

    /// Update RAG configuration
    pub fn set_config(&mut self, config: RagConfig) {
        self.vector_store.set_max_results(config.max_results);
        self.config = config;
    }

    /// Clear all documents
    pub fn clear(&mut self) {
        self.indexer.clear();
        self.vector_store.clear();
    }

    /// Save RAG system to a file using binary format
    ///
    /// Captures the complete state including vectors, full-text index, and configuration.
    pub fn save_to_file(&self, path: &str) -> Result<()> {
        let vector_store_data = self.vector_store.serialize()?;
        let indexer_data = self.indexer.serialize()?;
        let config_data = bincode::serialize(&self.config).map_err(|e| {
            ErrorWrapper::validation_failed(format!("Failed to serialize config: {}", e))
        })?;

        let snapshot = super::persistence::RagSystemSnapshot {
            version: super::persistence::Persistence::current_version(),
            vector_store_data,
            indexer_data,
            embeddings_model: self.embeddings.model_type().as_str().to_string(),
            config_data,
        };

        super::persistence::Persistence::save_binary(&snapshot, path)
    }

    /// Load RAG system from a file
    ///
    /// Restores vectors, full-text index, and configuration from saved state.
    pub fn load_from_file(path: &str) -> Result<Self> {
        let snapshot: super::persistence::RagSystemSnapshot =
            super::persistence::Persistence::load_binary(path)?;

        // Check version compatibility
        if !super::persistence::Persistence::is_version_compatible(snapshot.version) {
            return Err(ErrorWrapper::validation_failed(format!(
                "Incompatible snapshot version: {}",
                snapshot.version
            )));
        }

        // Deserialize components
        let vector_store =
            super::vector_store::VectorStore::deserialize(&snapshot.vector_store_data)?;
        let indexer = super::indexer::FullTextIndexer::deserialize(&snapshot.indexer_data)?;
        let config: RagConfig = bincode::deserialize(&snapshot.config_data).map_err(|e| {
            ErrorWrapper::validation_failed(format!("Failed to deserialize config: {}", e))
        })?;

        // Recreate embeddings service
        let model = EmbeddingModel::parse(&snapshot.embeddings_model);
        let embeddings = EmbeddingsService::new(model)?;

        Ok(RagSystem {
            embeddings,
            vector_store,
            indexer,
            config,
        })
    }
}

#[cfg(all(test, feature = "ai_backend"))]
mod tests {
    use super::*;

    #[test]
    fn test_rag_config_default() {
        let config = RagConfig::default();
        assert_eq!(config.semantic_weight, 0.6);
        assert_eq!(config.keyword_weight, 0.4);
        assert_eq!(config.max_results, 5);
    }

    #[test]
    fn test_rag_system_creation() {
        let config = RagConfig::default();
        let result = RagSystem::new(config);
        assert!(result.is_ok());
    }

    #[test]
    fn test_rag_add_documents_batch() {
        let mut rag = RagSystem::new(RagConfig::default()).unwrap();

        let docs = vec![
            ("doc1".to_string(), "hello world".to_string()),
            ("doc2".to_string(), "goodbye world".to_string()),
            ("doc3".to_string(), "testing batch operations".to_string()),
        ];

        assert!(rag.add_documents_batch(docs).is_ok());
        assert_eq!(rag.doc_count(), 3);

        // Verify all documents are retrievable
        let results = rag.retrieve("hello").unwrap();
        assert!(!results.is_empty());
    }

    #[test]
    fn test_rag_add_documents_batch_empty() {
        let mut rag = RagSystem::new(RagConfig::default()).unwrap();
        let docs = vec![];

        // Should handle empty batch gracefully
        assert!(rag.add_documents_batch(docs).is_ok());
        assert_eq!(rag.doc_count(), 0);
    }

    #[test]
    fn test_rag_remove_documents_batch() {
        let mut rag = RagSystem::new(RagConfig::default()).unwrap();

        // Add documents
        let docs = vec![
            ("doc1".to_string(), "hello world".to_string()),
            ("doc2".to_string(), "goodbye world".to_string()),
            ("doc3".to_string(), "testing batch operations".to_string()),
        ];
        rag.add_documents_batch(docs).unwrap();
        assert_eq!(rag.doc_count(), 3);

        // Remove batch
        let removed = rag.remove_documents_batch(&["doc1", "doc3"]);
        assert_eq!(removed, 2);
        assert_eq!(rag.doc_count(), 1);

        // Verify correct document remains
        let results = rag.retrieve("goodbye").unwrap();
        assert!(!results.is_empty());
    }

    #[test]
    fn test_rag_remove_documents_batch_partial() {
        let mut rag = RagSystem::new(RagConfig::default()).unwrap();

        // Add documents
        let docs = vec![
            ("doc1".to_string(), "hello world".to_string()),
            ("doc2".to_string(), "goodbye world".to_string()),
        ];
        rag.add_documents_batch(docs).unwrap();

        // Try to remove mix of existing and non-existing
        let removed = rag.remove_documents_batch(&["doc1", "doc3", "doc4"]);
        assert_eq!(removed, 1);
        assert_eq!(rag.doc_count(), 1);
    }

    #[test]
    fn test_rag_save_and_load() {
        use std::fs;
        use tempfile::NamedTempFile;

        let temp_file = NamedTempFile::new().unwrap();
        let path = temp_file.path().to_str().unwrap().to_string();

        // Create RAG system and add documents
        let mut rag = RagSystem::new(RagConfig::default()).unwrap();
        rag.add_document("doc1".to_string(), "hello world".to_string())
            .unwrap();
        rag.add_document("doc2".to_string(), "goodbye world".to_string())
            .unwrap();

        // Save
        rag.save_to_file(&path).unwrap();
        assert!(fs::metadata(&path).is_ok());

        // Load
        let mut loaded_rag = RagSystem::load_from_file(&path).unwrap();
        assert_eq!(loaded_rag.doc_count(), 2);

        // Verify search works on loaded system
        let results = loaded_rag.retrieve("hello").unwrap();
        assert!(!results.is_empty());

        fs::remove_file(&path).ok();
    }

    #[test]
    fn test_rag_batch_vs_sequential() {
        use std::time::Instant;

        // Sequential insertion
        let mut rag_seq = RagSystem::new(RagConfig::default()).unwrap();
        let docs: Vec<_> = (0..20)
            .map(|i| (format!("doc{}", i), format!("content number {}", i)))
            .collect();

        let start = Instant::now();
        assert!(rag_seq.add_documents(docs.clone()).is_ok());
        let seq_duration = start.elapsed();

        // Batch insertion
        let mut rag_batch = RagSystem::new(RagConfig::default()).unwrap();
        let start = Instant::now();
        assert!(rag_batch.add_documents_batch(docs).is_ok());
        let batch_duration = start.elapsed();

        // Both should have same document count
        assert_eq!(rag_seq.doc_count(), 20);
        assert_eq!(rag_batch.doc_count(), 20);

        // Batch should be faster or at least comparable
        // (batch avoids N cache rebuilds, seq rebuilds cache after each insert)
        let speedup = if batch_duration.as_nanos() > 0 {
            seq_duration.as_nanos() as f64 / batch_duration.as_nanos() as f64
        } else {
            1.0
        };
        println!(
            "Sequential: {:?}, Batch: {:?}, Speedup: {:.2}x",
            seq_duration, batch_duration, speedup
        );
    }

    #[test]
    fn test_rag_save_empty_system() {
        use std::fs;
        use tempfile::NamedTempFile;

        let temp_file = NamedTempFile::new().unwrap();
        let path = temp_file.path().to_str().unwrap().to_string();

        let rag = RagSystem::new(RagConfig::default()).unwrap();

        // Save empty system
        rag.save_to_file(&path).unwrap();

        // Load
        let loaded = RagSystem::load_from_file(&path).unwrap();
        assert_eq!(loaded.doc_count(), 0);

        fs::remove_file(&path).ok();
    }

    #[test]
    fn test_rag_load_nonexistent() {
        let result = RagSystem::load_from_file("/nonexistent/path.bin");
        assert!(result.is_err());
    }
}
