// BM25 Full-Text Search using Tantivy
// In-memory index for fast retrieval

use std::sync::Arc;

use parking_lot::RwLock;
use tantivy::collector::TopDocs;
use tantivy::query::QueryParser;
use tantivy::schema::{Schema, TextFieldIndexing, TextOptions, Value, TEXT};
use tantivy::{doc, Index, IndexWriter, ReloadPolicy};
use tracing::{debug, warn};

use crate::metrics::STORAGE_OPERATIONS;
use crate::storage::Chunk;
use crate::RLMError;

/// BM25 search result with score
#[derive(Debug, Clone)]
pub struct BM25Result {
    pub chunk_id: String,
    pub score: f32,
}

use tantivy::IndexReader;

/// BM25 index for full-text search
pub struct BM25Index {
    index: Arc<Index>,
    writer: Arc<RwLock<IndexWriter>>,
    reader: Arc<IndexReader>,
    schema: Schema,
}

impl BM25Index {
    /// Create a new in-memory BM25 index
    pub fn new() -> crate::Result<Self> {
        let mut schema_builder = Schema::builder();

        // chunk_id: TEXT + STORED (so we can retrieve it from search results)
        let chunk_id_options = TextOptions::default()
            .set_indexing_options(TextFieldIndexing::default())
            .set_stored();
        schema_builder.add_text_field("chunk_id", chunk_id_options);

        // content: TEXT (indexed for search, no need to store)
        schema_builder.add_text_field("content", TEXT);

        let schema = schema_builder.build();

        // Create in-memory index
        let index = Index::create_in_ram(schema.clone());

        // Create index writer (single writer, in-memory buffer)
        let writer = index
            .writer(50_000_000) // 50MB buffer
            .map_err(|e| RLMError::SearchError(format!("Failed to create index writer: {}", e)))?;

        // Create reader with OnCommit reload policy
        let reader = index
            .reader_builder()
            .reload_policy(ReloadPolicy::OnCommitWithDelay)
            .try_into()
            .map_err(|e| RLMError::SearchError(format!("Failed to create reader: {}", e)))?;

        debug!("Created in-memory BM25 index");

        Ok(Self {
            index: Arc::new(index),
            writer: Arc::new(RwLock::new(writer)),
            reader: Arc::new(reader),
            schema,
        })
    }

    /// Add a document to the index
    pub fn add_document(&self, chunk: &Chunk) -> crate::Result<()> {
        let chunk_id_field = self
            .schema
            .get_field("chunk_id")
            .map_err(|_| RLMError::SearchError("chunk_id field not found".to_string()))?;
        let content_field = self
            .schema
            .get_field("content")
            .map_err(|_| RLMError::SearchError("content field not found".to_string()))?;

        let doc = doc!(
            chunk_id_field => chunk.chunk_id.clone(),
            content_field => chunk.content.clone(),
        );

        let writer = self.writer.write();
        writer
            .add_document(doc)
            .map_err(|e| RLMError::SearchError(format!("Failed to add document: {}", e)))?;

        STORAGE_OPERATIONS
            .with_label_values(&["bm25_add_doc", "success"])
            .inc();

        Ok(())
    }

    /// Commit all pending documents
    pub fn commit(&self) -> crate::Result<()> {
        let mut writer = self.writer.write();
        writer
            .commit()
            .map_err(|e| RLMError::SearchError(format!("Failed to commit index: {}", e)))?;

        // Force reader reload to make committed documents visible
        self.reader
            .reload()
            .map_err(|e| RLMError::SearchError(format!("Failed to reload reader: {}", e)))?;

        debug!("Committed BM25 index and reloaded reader");
        Ok(())
    }

    /// Search the index using BM25
    pub fn search(&self, query_str: &str, limit: usize) -> crate::Result<Vec<BM25Result>> {
        // Use the existing reader (already has OnCommitWithDelay policy)
        // Force reload to ensure we see all committed documents
        self.reader
            .reload()
            .map_err(|e| RLMError::SearchError(format!("Failed to reload reader: {}", e)))?;

        let searcher = self.reader.searcher();

        debug!(
            "BM25 search: query='{}', index has {} docs",
            query_str,
            searcher.num_docs()
        );

        // Parse query
        let content_field = self
            .schema
            .get_field("content")
            .map_err(|_| RLMError::SearchError("content field not found".to_string()))?;

        let query_parser = QueryParser::for_index(&self.index, vec![content_field]);
        let query = query_parser
            .parse_query(query_str)
            .map_err(|e| RLMError::SearchError(format!("Failed to parse query: {}", e)))?;

        debug!("Parsed query: {:?}", query);

        // Search with BM25 scoring
        let top_docs = searcher
            .search(&query, &TopDocs::with_limit(limit))
            .map_err(|e| RLMError::SearchError(format!("Search failed: {}", e)))?;

        // Extract results
        let chunk_id_field = self
            .schema
            .get_field("chunk_id")
            .map_err(|_| RLMError::SearchError("chunk_id field not found".to_string()))?;

        let mut results = Vec::new();
        for (score, doc_address) in top_docs {
            let retrieved_doc = searcher
                .doc::<tantivy::TantivyDocument>(doc_address)
                .map_err(|e| RLMError::SearchError(format!("Failed to retrieve doc: {}", e)))?;

            // Get chunk_id as text
            let chunk_id_values: Vec<_> = retrieved_doc.get_all(chunk_id_field).collect();
            if let Some(first_value) = chunk_id_values.first() {
                if let Some(chunk_id_str) = first_value.as_str() {
                    results.push(BM25Result {
                        chunk_id: chunk_id_str.to_string(),
                        score,
                    });
                }
            }
        }

        STORAGE_OPERATIONS
            .with_label_values(&["bm25_search", "success"])
            .inc();

        debug!("BM25 search returned {} results", results.len());
        Ok(results)
    }

    /// Rebuild index from chunks
    pub fn rebuild_from_chunks(&self, chunks: &[Chunk]) -> crate::Result<()> {
        debug!("Rebuilding BM25 index from {} chunks", chunks.len());

        // Clear existing index
        {
            let mut writer = self.writer.write();
            writer
                .delete_all_documents()
                .map_err(|e| RLMError::SearchError(format!("Failed to clear index: {}", e)))?;
            writer
                .commit()
                .map_err(|e| RLMError::SearchError(format!("Failed to commit clear: {}", e)))?;
        }

        // Add all chunks
        for chunk in chunks {
            self.add_document(chunk)?;
        }

        // Commit
        self.commit()?;

        debug!("BM25 index rebuilt with {} documents", chunks.len());
        Ok(())
    }

    /// Get index statistics
    pub fn stats(&self) -> IndexStats {
        let reader = self
            .index
            .reader_builder()
            .reload_policy(ReloadPolicy::OnCommitWithDelay)
            .try_into();

        match reader {
            Ok(reader) => {
                let searcher = reader.searcher();
                IndexStats {
                    num_docs: searcher.num_docs() as usize,
                    num_segments: searcher.segment_readers().len(),
                }
            }
            Err(e) => {
                warn!("Failed to get index stats: {}", e);
                IndexStats {
                    num_docs: 0,
                    num_segments: 0,
                }
            }
        }
    }
}

impl Default for BM25Index {
    fn default() -> Self {
        Self::new().expect("Failed to create default BM25 index")
    }
}

/// Index statistics
#[derive(Debug, Clone)]
pub struct IndexStats {
    pub num_docs: usize,
    pub num_segments: usize,
}

#[cfg(test)]
mod tests {
    use chrono::Utc;

    use super::*;

    fn create_test_chunk(id: &str, content: &str) -> Chunk {
        Chunk {
            chunk_id: id.to_string(),
            doc_id: "test-doc".to_string(),
            content: content.to_string(),
            embedding: None,
            start_idx: 0,
            end_idx: content.len(),
            metadata: None,
            created_at: Utc::now().to_rfc3339(),
        }
    }

    #[test]
    fn test_bm25_index_creation() {
        let index = BM25Index::new();
        assert!(index.is_ok());
    }

    #[test]
    fn test_bm25_add_and_search() {
        let index = BM25Index::new().unwrap();

        // Add documents
        let chunk1 = create_test_chunk("chunk-1", "Rust programming language");
        let chunk2 = create_test_chunk("chunk-2", "Python programming tutorial");
        let chunk3 = create_test_chunk("chunk-3", "Rust async await patterns");

        index.add_document(&chunk1).unwrap();
        index.add_document(&chunk2).unwrap();
        index.add_document(&chunk3).unwrap();
        index.commit().unwrap();

        // Check stats
        let stats = index.stats();
        assert_eq!(stats.num_docs, 3, "Index should have 3 documents");

        // Search for "rust" (lowercase to match tokenization)
        let results = index.search("rust", 10).unwrap();
        assert!(
            results.len() >= 2,
            "Should find at least 2 results for 'rust', found {}",
            results.len()
        );

        // Search for "programming"
        let results = index.search("programming", 10).unwrap();
        assert!(
            results.len() >= 2,
            "Should find at least 2 results for 'programming', found {}",
            results.len()
        );
    }

    #[test]
    fn test_bm25_rebuild() {
        let index = BM25Index::new().unwrap();

        let chunks = vec![
            create_test_chunk("chunk-1", "First document"),
            create_test_chunk("chunk-2", "Second document"),
            create_test_chunk("chunk-3", "Third document"),
        ];

        index.rebuild_from_chunks(&chunks).unwrap();

        let results = index.search("document", 10).unwrap();
        assert_eq!(results.len(), 3);
    }

    #[test]
    fn test_bm25_stats() {
        let index = BM25Index::new().unwrap();

        let chunk = create_test_chunk("chunk-1", "Test content");
        index.add_document(&chunk).unwrap();
        index.commit().unwrap();

        let stats = index.stats();
        assert_eq!(stats.num_docs, 1);
    }

    #[test]
    fn test_bm25_relevance_ranking() {
        let index = BM25Index::new().unwrap();

        // Add documents with varying relevance
        let chunk1 = create_test_chunk("chunk-1", "error handling in Rust error error");
        let chunk2 = create_test_chunk("chunk-2", "Rust programming basics");
        let chunk3 = create_test_chunk("chunk-3", "error messages");

        index.add_document(&chunk1).unwrap();
        index.add_document(&chunk2).unwrap();
        index.add_document(&chunk3).unwrap();
        index.commit().unwrap();

        // Search for "error" - chunk-1 should rank highest (appears 3 times)
        let results = index.search("error", 10).unwrap();
        assert!(!results.is_empty());
        assert_eq!(results[0].chunk_id, "chunk-1"); // Most relevant
        assert!(results[0].score > 0.0);
    }
}
