//! Semantic (embedding-based) vector search.

use std::collections::HashMap;
use std::path::Path;

use usearch::{Index as UsearchIndex, IndexOptions, MetricKind, ScalarKind};

use crate::config::SearchConfig;
use crate::embeddings::{cosine_similarity, Embedder};
use crate::error::{Error, Result};
use crate::search::{normalize_scores, SearchQuery};
use crate::types::{Chunk, Document, MatchType, SearchResult};

/// Vector-based semantic search index.
pub struct SemanticIndex {
    index: UsearchIndex,
    dimension: usize,
    chunk_ids: Vec<String>,
    id_to_key: HashMap<String, u64>,
    #[allow(dead_code)]
    config: SearchConfig,
}

impl SemanticIndex {
    /// Create a new semantic index.
    pub fn new(dimension: usize, config: SearchConfig) -> Result<Self> {
        let options = IndexOptions {
            dimensions: dimension,
            metric: MetricKind::Cos, // Cosine similarity
            quantization: ScalarKind::F32,
            connectivity: 16,  // M parameter for HNSW
            expansion_add: 128, // efConstruction
            expansion_search: 64, // ef
            multi: false,
        };

        let index = UsearchIndex::new(&options)
            .map_err(|e| Error::index(format!("Failed to create vector index: {}", e)))?;

        Ok(Self {
            index,
            dimension,
            chunk_ids: Vec::new(),
            id_to_key: HashMap::new(),
            config,
        })
    }

    /// Load an existing index from disk.
    pub fn load(path: &Path, dimension: usize, config: SearchConfig) -> Result<Self> {
        let options = IndexOptions {
            dimensions: dimension,
            metric: MetricKind::Cos,
            quantization: ScalarKind::F32,
            connectivity: 16,
            expansion_add: 128,
            expansion_search: 64,
            multi: false,
        };

        let index = UsearchIndex::new(&options)
            .map_err(|e| Error::index(format!("Failed to create index: {}", e)))?;

        index
            .load(path.to_str().unwrap_or(""))
            .map_err(|e| Error::index(format!("Failed to load index: {}", e)))?;

        // Load chunk ID mapping
        let mapping_path = path.with_extension("mapping.json");
        let chunk_ids: Vec<String> = if mapping_path.exists() {
            let content = std::fs::read_to_string(&mapping_path)?;
            serde_json::from_str(&content)?
        } else {
            Vec::new()
        };

        let id_to_key: HashMap<String, u64> = chunk_ids
            .iter()
            .enumerate()
            .map(|(i, id)| (id.clone(), i as u64))
            .collect();

        Ok(Self {
            index,
            dimension,
            chunk_ids,
            id_to_key,
            config,
        })
    }

    /// Save the index to disk.
    pub fn save(&self, path: &Path) -> Result<()> {
        self.index
            .save(path.to_str().unwrap_or(""))
            .map_err(|e| Error::index(format!("Failed to save index: {}", e)))?;

        // Save chunk ID mapping
        let mapping_path = path.with_extension("mapping.json");
        let content = serde_json::to_string(&self.chunk_ids)?;
        std::fs::write(mapping_path, content)?;

        Ok(())
    }

    /// Add embeddings to the index.
    pub fn add_embeddings(&mut self, embeddings: &[(String, Vec<f32>)]) -> Result<()> {
        // Reserve capacity
        let new_size = self.chunk_ids.len() + embeddings.len();
        self.index
            .reserve(new_size)
            .map_err(|e| Error::index(format!("Failed to reserve capacity: {}", e)))?;

        for (chunk_id, embedding) in embeddings {
            if embedding.len() != self.dimension {
                return Err(Error::embedding(format!(
                    "Embedding dimension mismatch: expected {}, got {}",
                    self.dimension,
                    embedding.len()
                )));
            }

            let key = self.chunk_ids.len() as u64;
            self.index
                .add(key, embedding)
                .map_err(|e| Error::index(format!("Failed to add vector: {}", e)))?;

            self.id_to_key.insert(chunk_id.clone(), key);
            self.chunk_ids.push(chunk_id.clone());
        }

        Ok(())
    }

    /// Remove embeddings from the index.
    #[allow(dead_code)]
    pub fn remove_embeddings(&mut self, chunk_ids: &[String]) -> Result<()> {
        for chunk_id in chunk_ids {
            if let Some(&key) = self.id_to_key.get(chunk_id) {
                self.index
                    .remove(key)
                    .map_err(|e| Error::index(format!("Failed to remove vector: {}", e)))?;
            }
        }
        Ok(())
    }

    /// Clear the index by recreating it.
    #[allow(dead_code)]
    pub fn clear(&mut self) -> Result<()> {
        // Recreate the index since usearch doesn't have a clear method
        let options = IndexOptions {
            dimensions: self.dimension,
            metric: MetricKind::Cos,
            quantization: ScalarKind::F32,
            connectivity: 16,
            expansion_add: 128,
            expansion_search: 64,
            multi: false,
        };

        self.index = UsearchIndex::new(&options)
            .map_err(|e| Error::index(format!("Failed to recreate index: {}", e)))?;
        self.chunk_ids.clear();
        self.id_to_key.clear();
        Ok(())
    }

    /// Search for similar vectors.
    pub fn search_vectors(
        &self,
        query_embedding: &[f32],
        limit: usize,
    ) -> Result<Vec<(String, f32)>> {
        if query_embedding.len() != self.dimension {
            return Err(Error::embedding(format!(
                "Query dimension mismatch: expected {}, got {}",
                self.dimension,
                query_embedding.len()
            )));
        }

        let results = self
            .index
            .search(query_embedding, limit)
            .map_err(|e| Error::search(format!("Vector search failed: {}", e)))?;

        let matches: Vec<(String, f32)> = results
            .keys
            .into_iter()
            .zip(results.distances.into_iter())
            .filter_map(|(key, distance)| {
                self.chunk_ids
                    .get(key as usize)
                    .map(|id| {
                        // Convert distance to similarity score
                        // For cosine distance: similarity = 1 - distance
                        let similarity = 1.0 - distance;
                        (id.clone(), similarity)
                    })
            })
            .collect();

        Ok(matches)
    }

    /// Search using an embedder to encode the query.
    pub fn search_with_embedder<E: Embedder + ?Sized>(
        &self,
        query: &SearchQuery,
        embedder: &E,
        chunk_map: &HashMap<String, &Chunk>,
        doc_map: &HashMap<String, &Document>,
    ) -> Result<Vec<SearchResult>> {
        // Embed the query
        let query_embedding = embedder.embed(&query.text)?;

        // Search vectors
        let vector_results = self.search_vectors(&query_embedding, query.limit * 2)?;

        let mut results = Vec::new();

        for (chunk_id, score) in vector_results {
            if let Some(chunk) = chunk_map.get(&chunk_id) {
                let doc_id = chunk.document_id.0.to_string();
                if let Some(document) = doc_map.get(&doc_id) {
                    // Apply filter
                    if let Some(ref filter) = query.filter {
                        if !filter.matches(chunk, document) {
                            continue;
                        }
                    }

                    results.push(SearchResult {
                        chunk: (*chunk).clone(),
                        document: (*document).clone(),
                        score,
                        match_type: MatchType::Semantic,
                        highlights: vec![], // Semantic search doesn't have keyword highlights
                    });
                }
            }
        }

        // Normalize and filter
        normalize_scores(&mut results);
        results.retain(|r| r.score >= query.min_score);
        results.truncate(query.limit);

        Ok(results)
    }

    /// Get the number of vectors in the index.
    #[allow(dead_code)]
    pub fn len(&self) -> usize {
        self.chunk_ids.len()
    }

    /// Check if the index is empty.
    #[allow(dead_code)]
    pub fn is_empty(&self) -> bool {
        self.chunk_ids.is_empty()
    }
}

/// In-memory semantic search for smaller datasets.
#[allow(dead_code)]
pub struct InMemorySemanticSearch {
    embeddings: HashMap<String, Vec<f32>>,
    dimension: usize,
}

impl InMemorySemanticSearch {
    /// Create a new in-memory semantic search index.
    #[allow(dead_code)]
    pub fn new(dimension: usize) -> Self {
        Self {
            embeddings: HashMap::new(),
            dimension,
        }
    }

    /// Add an embedding to the index.
    #[allow(dead_code)]
    pub fn add(&mut self, chunk_id: String, embedding: Vec<f32>) {
        self.embeddings.insert(chunk_id, embedding);
    }

    /// Search for similar embeddings.
    #[allow(dead_code)]
    pub fn search(&self, query_embedding: &[f32], limit: usize) -> Vec<(String, f32)> {
        let mut scores: Vec<_> = self
            .embeddings
            .iter()
            .map(|(id, emb)| {
                let score = cosine_similarity(query_embedding, emb);
                (id.clone(), score)
            })
            .collect();

        scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        scores.truncate(limit);
        scores
    }
}

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

    #[test]
    fn test_in_memory_search() {
        let mut search = InMemorySemanticSearch::new(4);

        // Add some test embeddings
        search.add("chunk1".to_string(), vec![1.0, 0.0, 0.0, 0.0]);
        search.add("chunk2".to_string(), vec![0.0, 1.0, 0.0, 0.0]);
        search.add("chunk3".to_string(), vec![0.707, 0.707, 0.0, 0.0]);

        // Search for similar to [1, 0, 0, 0]
        let results = search.search(&[1.0, 0.0, 0.0, 0.0], 3);

        assert_eq!(results.len(), 3);
        assert_eq!(results[0].0, "chunk1"); // Exact match
        assert!((results[0].1 - 1.0).abs() < 0.01);
    }

    #[test]
    fn test_semantic_index_create() {
        let config = SearchConfig::default();
        let index = SemanticIndex::new(384, config).unwrap();
        assert!(index.is_empty());
    }
}
