// Semantic Search using Vector Similarity
// Cosine similarity-based ranking of chunks by embedding

use tracing::debug;

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

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

/// Semantic search using vector embeddings
pub struct SemanticSearch;

impl SemanticSearch {
    /// Create a new semantic search instance
    pub fn new() -> Self {
        Self
    }

    /// Search by embedding similarity
    pub fn search_by_embedding(
        chunks: &[Chunk],
        query_embedding: &[f32],
        limit: usize,
    ) -> Vec<SemanticResult> {
        debug!(
            "Semantic search with {} chunks, limit {}",
            chunks.len(),
            limit
        );

        // Filter chunks with embeddings and compute similarity
        let mut scored: Vec<(f32, String)> = chunks
            .iter()
            .filter_map(|chunk| {
                if let Some(ref embedding) = chunk.embedding {
                    let similarity = cosine_similarity(embedding, query_embedding);
                    Some((similarity, chunk.chunk_id.clone()))
                } else {
                    None
                }
            })
            .collect();

        // Sort by similarity descending
        scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));

        // Take top N
        let results: Vec<SemanticResult> = scored
            .into_iter()
            .take(limit)
            .map(|(similarity, chunk_id)| SemanticResult {
                chunk_id,
                similarity,
            })
            .collect();

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

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

    /// Rank chunks by similarity to query embedding
    pub fn rank_by_similarity(chunks: &[Chunk], query_embedding: &[f32]) -> Vec<(Chunk, f32)> {
        let mut scored: Vec<(Chunk, f32)> = chunks
            .iter()
            .filter_map(|chunk| {
                if let Some(ref embedding) = chunk.embedding {
                    let similarity = cosine_similarity(embedding, query_embedding);
                    Some((chunk.clone(), similarity))
                } else {
                    None
                }
            })
            .collect();

        // Sort by similarity descending
        scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        scored
    }
}

impl Default for SemanticSearch {
    fn default() -> Self {
        Self::new()
    }
}

/// Cosine similarity between two vectors
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    if a.len() != b.len() {
        return 0.0;
    }

    let dot_product: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let magnitude_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let magnitude_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();

    if magnitude_a == 0.0 || magnitude_b == 0.0 {
        return 0.0;
    }

    dot_product / (magnitude_a * magnitude_b)
}

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

    use super::*;

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

    #[test]
    fn test_cosine_similarity() {
        // Identical vectors
        let a = vec![1.0, 0.0, 0.0];
        let b = vec![1.0, 0.0, 0.0];
        assert!((cosine_similarity(&a, &b) - 1.0).abs() < 0.001);

        // Orthogonal vectors
        let c = vec![1.0, 0.0, 0.0];
        let d = vec![0.0, 1.0, 0.0];
        assert!((cosine_similarity(&c, &d) - 0.0).abs() < 0.001);

        // Similar vectors
        let e = vec![1.0, 1.0, 0.0];
        let f = vec![1.0, 0.0, 0.0];
        let similarity = cosine_similarity(&e, &f);
        assert!(similarity > 0.7 && similarity < 0.8);
    }

    #[test]
    fn test_cosine_similarity_edge_cases() {
        // Empty vectors
        assert_eq!(cosine_similarity(&[], &[]), 0.0);

        // Different lengths
        assert_eq!(cosine_similarity(&[1.0], &[1.0, 2.0]), 0.0);

        // Zero vectors
        assert_eq!(cosine_similarity(&[0.0, 0.0], &[1.0, 1.0]), 0.0);
    }

    #[test]
    fn test_semantic_search() {
        let chunks = vec![
            create_test_chunk_with_embedding("chunk-1", "content1", vec![1.0, 0.0, 0.0]),
            create_test_chunk_with_embedding("chunk-2", "content2", vec![0.0, 1.0, 0.0]),
            create_test_chunk_with_embedding("chunk-3", "content3", vec![0.9, 0.1, 0.0]),
        ];

        let query_embedding = vec![1.0, 0.0, 0.0];
        let results = SemanticSearch::search_by_embedding(&chunks, &query_embedding, 2);

        assert_eq!(results.len(), 2);
        // chunk-1 should be first (exact match)
        assert_eq!(results[0].chunk_id, "chunk-1");
        assert!(results[0].similarity > 0.99);
        // chunk-3 should be second (similar)
        assert_eq!(results[1].chunk_id, "chunk-3");
    }

    #[test]
    fn test_semantic_search_no_embeddings() {
        let chunks = vec![Chunk {
            chunk_id: "chunk-1".to_string(),
            doc_id: "test".to_string(),
            content: "test".to_string(),
            embedding: None, // No embedding
            start_idx: 0,
            end_idx: 4,
            metadata: None,
            created_at: Utc::now().to_rfc3339(),
        }];

        let query_embedding = vec![1.0, 0.0, 0.0];
        let results = SemanticSearch::search_by_embedding(&chunks, &query_embedding, 10);

        assert_eq!(results.len(), 0); // Should skip chunks without embeddings
    }

    #[test]
    fn test_rank_by_similarity() {
        let chunks = vec![
            create_test_chunk_with_embedding("chunk-1", "content1", vec![1.0, 0.0, 0.0]),
            create_test_chunk_with_embedding("chunk-2", "content2", vec![0.0, 1.0, 0.0]),
            create_test_chunk_with_embedding("chunk-3", "content3", vec![0.9, 0.1, 0.0]),
        ];

        let query_embedding = vec![1.0, 0.0, 0.0];
        let ranked = SemanticSearch::rank_by_similarity(&chunks, &query_embedding);

        assert_eq!(ranked.len(), 3);
        // Should be sorted by similarity
        assert_eq!(ranked[0].0.chunk_id, "chunk-1");
        assert!(ranked[0].1 > ranked[1].1);
        assert!(ranked[1].1 > ranked[2].1);
    }
}
