// Hybrid Search: BM25 + Semantic + RRF Fusion
// Combines keyword search and vector similarity for optimal retrieval

use std::sync::Arc;
use std::time::Instant;

use tracing::debug;

use super::bm25::{BM25Index, BM25Result};
use super::rrf::{reciprocal_rank_fusion_scored, RRFConfig};
use super::semantic::{SemanticResult, SemanticSearch};
use crate::metrics::{QUERY_DURATION, STORAGE_OPERATIONS};
use crate::storage::{Chunk, Storage};

/// Scored chunk from hybrid search
#[derive(Debug, Clone)]
pub struct ScoredChunk {
    pub chunk: Chunk,
    pub score: f32,
    pub bm25_score: Option<f32>,
    pub semantic_score: Option<f32>,
}

/// Hybrid search configuration
#[derive(Debug, Clone)]
pub struct HybridSearchConfig {
    /// RRF configuration
    pub rrf_config: RRFConfig,
    /// Weight for BM25 results (0.0 - 1.0)
    pub bm25_weight: f32,
    /// Weight for semantic results (0.0 - 1.0)
    pub semantic_weight: f32,
}

impl Default for HybridSearchConfig {
    fn default() -> Self {
        Self {
            rrf_config: RRFConfig::default(),
            bm25_weight: 0.5,
            semantic_weight: 0.5,
        }
    }
}

/// Hybrid search orchestrator
pub struct HybridSearch<S: Storage> {
    storage: Arc<S>,
    bm25_index: Arc<BM25Index>,
    #[allow(dead_code)]
    semantic_search: SemanticSearch,
    config: HybridSearchConfig,
}

impl<S: Storage> HybridSearch<S> {
    /// Create a new hybrid search instance
    pub fn new(storage: Arc<S>, bm25_index: Arc<BM25Index>) -> crate::Result<Self> {
        Ok(Self {
            storage,
            bm25_index,
            semantic_search: SemanticSearch::new(),
            config: HybridSearchConfig::default(),
        })
    }

    /// Create with custom configuration
    pub fn with_config(
        storage: Arc<S>,
        bm25_index: Arc<BM25Index>,
        config: HybridSearchConfig,
    ) -> crate::Result<Self> {
        Ok(Self {
            storage,
            bm25_index,
            semantic_search: SemanticSearch::new(),
            config,
        })
    }

    /// Hybrid search: BM25 + Semantic + RRF fusion
    ///
    /// # Arguments
    /// - `doc_id`: Document to search within
    /// - `query_text`: Keyword query for BM25
    /// - `query_embedding`: Vector embedding for semantic search
    /// - `limit`: Maximum results to return
    pub async fn search(
        &self,
        doc_id: &str,
        query_text: &str,
        query_embedding: &[f32],
        limit: usize,
    ) -> crate::Result<Vec<ScoredChunk>> {
        let start = Instant::now();

        debug!(
            "Hybrid search: doc={}, query='{}', limit={}",
            doc_id, query_text, limit
        );

        // Get all chunks for the document
        let chunks = self.storage.get_chunks(doc_id).await?;

        if chunks.is_empty() {
            debug!("No chunks found for doc {}", doc_id);
            return Ok(Vec::new());
        }

        // Run BM25 and semantic search in parallel
        let bm25_results = self.bm25_search(query_text, limit * 2)?;
        let semantic_results =
            SemanticSearch::search_by_embedding(&chunks, query_embedding, limit * 2);

        // Prepare ranked lists for RRF
        let bm25_ranked: Vec<(String, f32)> = bm25_results
            .iter()
            .map(|r| (r.chunk_id.clone(), r.score))
            .collect();

        let semantic_ranked: Vec<(String, f32)> = semantic_results
            .iter()
            .map(|r| (r.chunk_id.clone(), r.similarity))
            .collect();

        // Apply RRF fusion
        let fused_results = reciprocal_rank_fusion_scored(
            &[bm25_ranked, semantic_ranked],
            self.config.rrf_config.clone(),
            limit,
        );

        // Map back to chunks with scores
        let mut scored_chunks = Vec::new();
        for rrf_result in fused_results {
            if let Some(chunk) = chunks.iter().find(|c| c.chunk_id == rrf_result.chunk_id) {
                // Find original scores
                let bm25_score = bm25_results
                    .iter()
                    .find(|r| r.chunk_id == rrf_result.chunk_id)
                    .map(|r| r.score);

                let semantic_score = semantic_results
                    .iter()
                    .find(|r| r.chunk_id == rrf_result.chunk_id)
                    .map(|r| r.similarity);

                scored_chunks.push(ScoredChunk {
                    chunk: chunk.clone(),
                    score: rrf_result.score,
                    bm25_score,
                    semantic_score,
                });
            }
        }

        let duration = start.elapsed();
        QUERY_DURATION
            .with_label_values(&["hybrid"])
            .observe(duration.as_secs_f64());

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

        debug!(
            "Hybrid search completed in {:?}, returned {} results",
            duration,
            scored_chunks.len()
        );

        Ok(scored_chunks)
    }

    /// BM25-only search
    pub fn bm25_search(&self, query_text: &str, limit: usize) -> crate::Result<Vec<BM25Result>> {
        let start = Instant::now();
        let results = self.bm25_index.search(query_text, limit)?;

        let duration = start.elapsed();
        QUERY_DURATION
            .with_label_values(&["bm25"])
            .observe(duration.as_secs_f64());

        Ok(results)
    }

    /// Semantic-only search
    pub async fn semantic_search(
        &self,
        doc_id: &str,
        query_embedding: &[f32],
        limit: usize,
    ) -> crate::Result<Vec<SemanticResult>> {
        let start = Instant::now();

        let chunks = self.storage.get_chunks(doc_id).await?;
        let results = SemanticSearch::search_by_embedding(&chunks, query_embedding, limit);

        let duration = start.elapsed();
        QUERY_DURATION
            .with_label_values(&["semantic"])
            .observe(duration.as_secs_f64());

        Ok(results)
    }

    /// Rebuild BM25 index from document chunks
    pub async fn rebuild_index(&self, doc_id: &str) -> crate::Result<()> {
        let chunks = self.storage.get_chunks(doc_id).await?;
        self.bm25_index.rebuild_from_chunks(&chunks)?;
        Ok(())
    }

    /// Get index statistics
    pub fn index_stats(&self) -> super::bm25::IndexStats {
        self.bm25_index.stats()
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::sync::Mutex;

    use async_trait::async_trait;
    use chrono::Utc;

    use super::*;
    use crate::storage::{Chunk, ExecutionHistory, Storage};

    // Mock storage for testing
    struct MockStorage {
        chunks: Arc<Mutex<HashMap<String, Vec<Chunk>>>>,
    }

    impl MockStorage {
        fn new() -> Self {
            Self {
                chunks: Arc::new(Mutex::new(HashMap::new())),
            }
        }

        fn add_chunk(&self, chunk: Chunk) {
            let mut chunks = self.chunks.lock().unwrap();
            chunks.entry(chunk.doc_id.clone()).or_default().push(chunk);
        }
    }

    #[async_trait]
    impl Storage for MockStorage {
        async fn save_chunk(&self, chunk: Chunk) -> crate::Result<()> {
            self.add_chunk(chunk);
            Ok(())
        }

        async fn get_chunks(&self, doc_id: &str) -> crate::Result<Vec<Chunk>> {
            let chunks = self.chunks.lock().unwrap();
            Ok(chunks.get(doc_id).cloned().unwrap_or_default())
        }

        async fn get_chunk(&self, _chunk_id: &str) -> crate::Result<Option<Chunk>> {
            Ok(None)
        }

        async fn search_by_embedding(
            &self,
            _embedding: &[f32],
            _limit: usize,
        ) -> crate::Result<Vec<Chunk>> {
            Ok(Vec::new())
        }

        async fn save_buffer(&self, _buffer: crate::storage::Buffer) -> crate::Result<()> {
            Ok(())
        }

        async fn get_buffer(
            &self,
            _buffer_id: &str,
        ) -> crate::Result<Option<crate::storage::Buffer>> {
            Ok(None)
        }

        async fn cleanup_expired_buffers(&self) -> crate::Result<u64> {
            Ok(0)
        }

        async fn save_execution(&self, _execution: ExecutionHistory) -> crate::Result<()> {
            Ok(())
        }

        async fn get_executions(
            &self,
            _doc_id: &str,
            _limit: usize,
        ) -> crate::Result<Vec<ExecutionHistory>> {
            Ok(Vec::new())
        }

        async fn delete_chunks(&self, _doc_id: &str) -> crate::Result<u64> {
            Ok(0)
        }
    }

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

    #[tokio::test]
    async fn test_hybrid_search_basic() {
        let storage = Arc::new(MockStorage::new());
        let bm25_index = Arc::new(BM25Index::new().unwrap());

        // Add test chunks
        let chunk1 = create_test_chunk(
            "chunk-1",
            "doc-1",
            "Rust programming language",
            vec![1.0, 0.0, 0.0],
        );
        let chunk2 = create_test_chunk(
            "chunk-2",
            "doc-1",
            "Python programming tutorial",
            vec![0.0, 1.0, 0.0],
        );
        let chunk3 = create_test_chunk(
            "chunk-3",
            "doc-1",
            "Rust async patterns",
            vec![0.9, 0.1, 0.0],
        );

        storage.add_chunk(chunk1.clone());
        storage.add_chunk(chunk2.clone());
        storage.add_chunk(chunk3.clone());

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

        let hybrid = HybridSearch::new(storage, bm25_index).unwrap();

        // Search for "Rust" with embedding similar to chunk1
        let query_embedding = vec![1.0, 0.0, 0.0];
        let results = hybrid
            .search("doc-1", "Rust", &query_embedding, 2)
            .await
            .unwrap();

        assert!(!results.is_empty());
        // Should return chunk-1 and/or chunk-3 (both match "Rust" and have similar
        // embeddings)
        assert!(results.iter().any(|r| r.chunk.chunk_id == "chunk-1"));
    }

    #[tokio::test]
    async fn test_hybrid_search_empty_doc() {
        let storage = Arc::new(MockStorage::new());
        let bm25_index = Arc::new(BM25Index::new().unwrap());

        let hybrid = HybridSearch::new(storage, bm25_index).unwrap();

        let results = hybrid
            .search("nonexistent-doc", "query", &[1.0, 0.0], 10)
            .await
            .unwrap();

        assert!(results.is_empty());
    }

    #[tokio::test]
    async fn test_bm25_only_search() {
        let storage = Arc::new(MockStorage::new());
        let bm25_index = Arc::new(BM25Index::new().unwrap());

        let chunk = create_test_chunk("chunk-1", "doc-1", "test content", vec![1.0, 0.0]);
        bm25_index.add_document(&chunk).unwrap();
        bm25_index.commit().unwrap();

        let hybrid = HybridSearch::new(storage, bm25_index).unwrap();

        let results = hybrid.bm25_search("test", 10).unwrap();
        assert!(!results.is_empty());
    }

    #[tokio::test]
    async fn test_semantic_only_search() {
        let storage = Arc::new(MockStorage::new());
        let bm25_index = Arc::new(BM25Index::new().unwrap());

        let chunk = create_test_chunk("chunk-1", "doc-1", "test", vec![1.0, 0.0, 0.0]);
        storage.add_chunk(chunk);

        let hybrid = HybridSearch::new(storage, bm25_index).unwrap();

        let results = hybrid
            .semantic_search("doc-1", &[1.0, 0.0, 0.0], 10)
            .await
            .unwrap();
        assert!(!results.is_empty());
    }
}
