use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::debug;

/// Vector database trait for embeddings storage and retrieval
#[async_trait::async_trait]
pub trait VectorDB: Send + Sync {
    /// Insert a vector with metadata
    async fn insert(&self, id: String, vector: Vec<f32>, metadata: serde_json::Value) -> Result<(), VectorError>;
    
    /// Search for similar vectors
    async fn search(&self, query: Vec<f32>, top_k: usize) -> Result<Vec<SearchResult>, VectorError>;
    
    /// Delete a vector by ID
    async fn delete(&self, id: &str) -> Result<(), VectorError>;
    
    /// Get vector by ID
    async fn get(&self, id: &str) -> Option<(Vec<f32>, serde_json::Value)>;
}

/// Search result
#[derive(Debug, Clone)]
pub struct SearchResult {
    pub id: String,
    pub score: f32,
    pub metadata: serde_json::Value,
}

/// Vector operations errors
#[derive(Debug, thiserror::Error)]
pub enum VectorError {
    #[error("Vector operation failed: {0}")]
    OperationError(String),
    
    #[error("Invalid vector dimension")]
    InvalidDimension,
    
    #[error("Vector not found")]
    NotFound,
}

/// Simple in-memory vector store (for MVP)
pub struct InMemoryVectorDB {
    vectors: Arc<RwLock<std::collections::HashMap<String, (Vec<f32>, serde_json::Value)>>>,
    dimension: usize,
}

impl InMemoryVectorDB {
    pub fn new(dimension: usize) -> Self {
        Self {
            vectors: Arc::new(RwLock::new(std::collections::HashMap::new())),
            dimension,
        }
    }
    
    /// Compute cosine similarity
    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 norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
        let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
        
        if norm_a == 0.0 || norm_b == 0.0 {
            return 0.0;
        }
        
        dot_product / (norm_a * norm_b)
    }
}

#[async_trait::async_trait]
impl VectorDB for InMemoryVectorDB {
    async fn insert(&self, id: String, vector: Vec<f32>, metadata: serde_json::Value) -> Result<(), VectorError> {
        if vector.len() != self.dimension {
            return Err(VectorError::InvalidDimension);
        }
        
        let mut vectors = self.vectors.write().await;
        vectors.insert(id.clone(), (vector, metadata));
        debug!("Inserted vector: {}", id);
        Ok(())
    }
    
    async fn search(&self, query: Vec<f32>, top_k: usize) -> Result<Vec<SearchResult>, VectorError> {
        if query.len() != self.dimension {
            return Err(VectorError::InvalidDimension);
        }
        
        let vectors = self.vectors.read().await;
        let mut results: Vec<SearchResult> = vectors
            .iter()
            .map(|(id, (vector, metadata))| {
                let score = Self::cosine_similarity(&query, vector);
                SearchResult {
                    id: id.clone(),
                    score,
                    metadata: metadata.clone(),
                }
            })
            .collect();
        
        // Sort by score (descending) and take top_k
        results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
        results.truncate(top_k);
        
        debug!("Vector search returned {} results", results.len());
        Ok(results)
    }
    
    async fn delete(&self, id: &str) -> Result<(), VectorError> {
        let mut vectors = self.vectors.write().await;
        vectors.remove(id);
        debug!("Deleted vector: {}", id);
        Ok(())
    }
    
    async fn get(&self, id: &str) -> Option<(Vec<f32>, serde_json::Value)> {
        let vectors = self.vectors.read().await;
        vectors.get(id).cloned()
    }
}

