litellm-rs 0.1.1

//! Vector database implementation
//!
//! This module provides vector storage and similarity search functionality.

use crate::config::VectorDbConfig;
use crate::utils::error::{GatewayError, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use tracing::{debug, info};

/// Vector data for storage and retrieval
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VectorData {
    /// Unique identifier
    pub id: String,
    /// Vector embedding
    pub vector: Vec<f32>,
    /// Associated metadata
    pub metadata: HashMap<String, serde_json::Value>,
}

/// Vector store trait
#[async_trait::async_trait]
pub trait VectorStore: Send + Sync {
    /// Search for similar vectors
    async fn search(&self, vector: Vec<f32>, limit: usize) -> Result<Vec<SearchResult>>;

    /// Insert vectors
    async fn insert(&self, vectors: Vec<VectorData>) -> Result<()>;

    /// Delete vectors by ID
    async fn delete(&self, ids: Vec<String>) -> Result<()>;
}

/// Vector store backend enum
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub enum VectorStoreBackend {
    /// Qdrant vector database
    Qdrant(QdrantStore),
    /// Weaviate vector database
    Weaviate(WeaviateStore),
    /// Pinecone vector database
    Pinecone(PineconeStore),
}

/// Qdrant vector store
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct QdrantStore {
    url: String,
    api_key: Option<String>,
    collection: String,
    client: reqwest::Client,
}

/// Weaviate vector store
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct WeaviateStore {
    url: String,
    api_key: Option<String>,
    collection: String,
    client: reqwest::Client,
}

/// Pinecone vector store
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct PineconeStore {
    url: String,
    api_key: Option<String>,
    collection: String,
    client: reqwest::Client,
}

/// Search result from vector database
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchResult {
    /// Result ID
    pub id: String,
    /// Similarity score
    pub score: f32,
    /// Associated metadata
    pub metadata: Option<serde_json::Value>,
    /// Vector data (optional)
    pub vector: Option<Vec<f32>>,
}

/// Vector point for storage
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VectorPoint {
    /// Point ID
    pub id: String,
    /// Vector data
    pub vector: Vec<f32>,
    /// Associated metadata
    pub metadata: Option<serde_json::Value>,
}

/// Pinecone vector store implementation
#[allow(dead_code)]
pub struct PineconeVectorStore {
    config: VectorDbConfig,
    client: reqwest::Client,
}

#[allow(dead_code)]
impl VectorStoreBackend {
    /// Create a new vector store instance
    pub async fn new(config: &VectorDbConfig) -> Result<Self> {
        info!("Initializing vector database: {}", config.db_type);

        match config.db_type.as_str() {
            "qdrant" => Ok(VectorStoreBackend::Qdrant(QdrantStore::new(config).await?)),
            "weaviate" => Ok(VectorStoreBackend::Weaviate(
                WeaviateStore::new(config).await?,
            )),
            "pinecone" => Ok(VectorStoreBackend::Pinecone(
                PineconeStore::new(config).await?,
            )),
            _ => Err(GatewayError::Config(format!(
                "Unsupported vector DB type: {}",
                config.db_type
            ))),
        }
    }

    /// Store a vector with metadata
    pub async fn store(
        &self,
        id: &str,
        vector: &[f32],
        metadata: Option<serde_json::Value>,
    ) -> Result<()> {
        match self {
            VectorStoreBackend::Qdrant(store) => store.store(id, vector, metadata).await,
            VectorStoreBackend::Weaviate(store) => store.store(id, vector, metadata).await,
            VectorStoreBackend::Pinecone(store) => store.store(id, vector, metadata).await,
        }
    }

    /// Search for similar vectors
    pub async fn search(
        &self,
        query_vector: &[f32],
        limit: usize,
        threshold: Option<f32>,
    ) -> Result<Vec<SearchResult>> {
        match self {
            VectorStoreBackend::Qdrant(store) => store.search(query_vector, limit, threshold).await,
            VectorStoreBackend::Weaviate(store) => {
                store.search(query_vector, limit, threshold).await
            }
            VectorStoreBackend::Pinecone(store) => {
                store.search(query_vector, limit, threshold).await
            }
        }
    }

    /// Delete a vector by ID
    pub async fn delete(&self, id: &str) -> Result<()> {
        match self {
            VectorStoreBackend::Qdrant(store) => store.delete(id).await,
            VectorStoreBackend::Weaviate(store) => store.delete(id).await,
            VectorStoreBackend::Pinecone(store) => store.delete(id).await,
        }
    }

    /// Get a vector by ID
    pub async fn get(&self, id: &str) -> Result<Option<VectorPoint>> {
        match self {
            VectorStoreBackend::Qdrant(store) => store.get(id).await,
            VectorStoreBackend::Weaviate(store) => store.get(id).await,
            VectorStoreBackend::Pinecone(store) => store.get(id).await,
        }
    }

    /// Health check
    pub async fn health_check(&self) -> Result<()> {
        match self {
            VectorStoreBackend::Qdrant(store) => store.health_check().await,
            VectorStoreBackend::Weaviate(store) => store.health_check().await,
            VectorStoreBackend::Pinecone(store) => store.health_check().await,
        }
    }

    /// Close connections
    pub async fn close(&self) -> Result<()> {
        match self {
            VectorStoreBackend::Qdrant(_store) => Ok(()), // No explicit close needed for HTTP clients
            VectorStoreBackend::Weaviate(_store) => Ok(()),
            VectorStoreBackend::Pinecone(_store) => Ok(()),
        }
    }

    /// Batch store vectors
    pub async fn batch_store(&self, points: &[VectorPoint]) -> Result<()> {
        match self {
            VectorStoreBackend::Qdrant(store) => store.batch_store(points).await,
            VectorStoreBackend::Weaviate(store) => store.batch_store(points).await,
            VectorStoreBackend::Pinecone(store) => store.batch_store(points).await,
        }
    }

    /// Count vectors in collection
    pub async fn count(&self) -> Result<u64> {
        match self {
            VectorStoreBackend::Qdrant(store) => store.count().await,
            VectorStoreBackend::Weaviate(store) => store.count().await,
            VectorStoreBackend::Pinecone(store) => store.count().await,
        }
    }
}

#[allow(dead_code)]
impl QdrantStore {
    /// Create a new Qdrant store
    pub async fn new(config: &VectorDbConfig) -> Result<Self> {
        let client = reqwest::Client::new();

        let store = Self {
            url: config.url.clone(),
            api_key: Some(config.api_key.clone()),
            collection: config.index_name.clone(),
            client,
        };

        // Ensure collection exists
        store.ensure_collection().await?;

        info!("Qdrant vector store initialized");
        Ok(store)
    }

    /// Ensure collection exists
    async fn ensure_collection(&self) -> Result<()> {
        let url = format!("{}/collections/{}", self.url, self.collection);
        let mut request = self.client.get(&url);

        if let Some(api_key) = &self.api_key {
            request = request.header("api-key", api_key);
        }

        let response = request
            .send()
            .await
            .map_err(|e| GatewayError::VectorDb(format!("Failed to check collection: {}", e)))?;

        if response.status() == 404 {
            // Collection doesn't exist, create it
            self.create_collection().await?;
        } else if !response.status().is_success() {
            return Err(GatewayError::VectorDb(format!(
                "Failed to check collection: {}",
                response.status()
            )));
        }

        Ok(())
    }

    /// Create collection
    async fn create_collection(&self) -> Result<()> {
        let url = format!("{}/collections/{}", self.url, self.collection);
        let payload = serde_json::json!({
            "vectors": {
                "size": 1536, // Default OpenAI embedding size
                "distance": "Cosine"
            }
        });

        let mut request = self.client.put(&url).json(&payload);

        if let Some(api_key) = &self.api_key {
            request = request.header("api-key", api_key);
        }

        let response = request
            .send()
            .await
            .map_err(|e| GatewayError::VectorDb(format!("Failed to create collection: {}", e)))?;

        if !response.status().is_success() {
            return Err(GatewayError::VectorDb(format!(
                "Failed to create collection: {}",
                response.status()
            )));
        }

        info!("Created Qdrant collection: {}", self.collection);
        Ok(())
    }

    /// Store a vector
    pub async fn store(
        &self,
        id: &str,
        vector: &[f32],
        metadata: Option<serde_json::Value>,
    ) -> Result<()> {
        let url = format!("{}/collections/{}/points", self.url, self.collection);
        let payload = serde_json::json!({
            "points": [{
                "id": id,
                "vector": vector,
                "payload": metadata.unwrap_or_default()
            }]
        });

        let mut request = self.client.put(&url).json(&payload);

        if let Some(api_key) = &self.api_key {
            request = request.header("api-key", api_key);
        }

        let response = request
            .send()
            .await
            .map_err(|e| GatewayError::VectorDb(format!("Failed to store vector: {}", e)))?;

        if !response.status().is_success() {
            return Err(GatewayError::VectorDb(format!(
                "Failed to store vector: {}",
                response.status()
            )));
        }

        debug!("Stored vector: {}", id);
        Ok(())
    }

    /// Search for similar vectors
    pub async fn search(
        &self,
        query_vector: &[f32],
        limit: usize,
        threshold: Option<f32>,
    ) -> Result<Vec<SearchResult>> {
        let url = format!("{}/collections/{}/points/search", self.url, self.collection);
        let mut payload = serde_json::json!({
            "vector": query_vector,
            "limit": limit,
            "with_payload": true,
            "with_vector": false
        });

        if let Some(threshold) = threshold {
            payload["score_threshold"] =
                serde_json::Value::Number(serde_json::Number::from_f64(threshold as f64).unwrap());
        }

        let mut request = self.client.post(&url).json(&payload);

        if let Some(api_key) = &self.api_key {
            request = request.header("api-key", api_key);
        }

        let response = request
            .send()
            .await
            .map_err(|e| GatewayError::VectorDb(format!("Failed to search vectors: {}", e)))?;

        if !response.status().is_success() {
            return Err(GatewayError::VectorDb(format!(
                "Failed to search vectors: {}",
                response.status()
            )));
        }

        let result: serde_json::Value = response.json().await.map_err(|e| {
            GatewayError::VectorDb(format!("Failed to parse search response: {}", e))
        })?;

        let mut search_results = Vec::new();
        if let Some(points) = result["result"].as_array() {
            for point in points {
                if let (Some(id), Some(score)) = (point["id"].as_str(), point["score"].as_f64()) {
                    search_results.push(SearchResult {
                        id: id.to_string(),
                        score: score as f32,
                        metadata: point["payload"].clone().into(),
                        vector: None,
                    });
                }
            }
        }

        Ok(search_results)
    }

    /// Delete a vector
    pub async fn delete(&self, id: &str) -> Result<()> {
        let url = format!("{}/collections/{}/points/delete", self.url, self.collection);
        let payload = serde_json::json!({
            "points": [id]
        });

        let mut request = self.client.post(&url).json(&payload);

        if let Some(api_key) = &self.api_key {
            request = request.header("api-key", api_key);
        }

        let response = request
            .send()
            .await
            .map_err(|e| GatewayError::VectorDb(format!("Failed to delete vector: {}", e)))?;

        if !response.status().is_success() {
            return Err(GatewayError::VectorDb(format!(
                "Failed to delete vector: {}",
                response.status()
            )));
        }

        debug!("Deleted vector: {}", id);
        Ok(())
    }

    /// Get a vector by ID
    pub async fn get(&self, id: &str) -> Result<Option<VectorPoint>> {
        let url = format!("{}/collections/{}/points/{}", self.url, self.collection, id);
        let mut request = self.client.get(&url);

        if let Some(api_key) = &self.api_key {
            request = request.header("api-key", api_key);
        }

        let response = request
            .send()
            .await
            .map_err(|e| GatewayError::VectorDb(format!("Failed to get vector: {}", e)))?;

        if response.status() == 404 {
            return Ok(None);
        }

        if !response.status().is_success() {
            return Err(GatewayError::VectorDb(format!(
                "Failed to get vector: {}",
                response.status()
            )));
        }

        let result: serde_json::Value = response
            .json()
            .await
            .map_err(|e| GatewayError::VectorDb(format!("Failed to parse get response: {}", e)))?;

        if let Some(point) = result["result"].as_object() {
            if let (Some(id), Some(vector)) = (point["id"].as_str(), point["vector"].as_array()) {
                let vector_data: Vec<f32> = vector
                    .iter()
                    .filter_map(|v| v.as_f64().map(|f| f as f32))
                    .collect();

                return Ok(Some(VectorPoint {
                    id: id.to_string(),
                    vector: vector_data,
                    metadata: point["payload"].clone().into(),
                }));
            }
        }

        Ok(None)
    }

    /// Health check
    pub async fn health_check(&self) -> Result<()> {
        let url = format!("{}/", self.url);
        let mut request = self.client.get(&url);

        if let Some(api_key) = &self.api_key {
            request = request.header("api-key", api_key);
        }

        let response = request
            .send()
            .await
            .map_err(|e| GatewayError::VectorDb(format!("Qdrant health check failed: {}", e)))?;

        if !response.status().is_success() {
            return Err(GatewayError::VectorDb(format!(
                "Qdrant health check failed: {}",
                response.status()
            )));
        }

        Ok(())
    }

    /// Close connections
    pub async fn close(&self) -> Result<()> {
        // HTTP client doesn't need explicit closing
        Ok(())
    }

    /// Batch store vectors
    pub async fn batch_store(&self, points: &[VectorPoint]) -> Result<()> {
        let url = format!("{}/collections/{}/points", self.url, self.collection);
        let qdrant_points: Vec<serde_json::Value> = points
            .iter()
            .map(|point| {
                serde_json::json!({
                    "id": point.id,
                    "vector": point.vector,
                    "payload": point.metadata.clone().unwrap_or_default()
                })
            })
            .collect();

        let payload = serde_json::json!({
            "points": qdrant_points
        });

        let mut request = self.client.put(&url).json(&payload);

        if let Some(api_key) = &self.api_key {
            request = request.header("api-key", api_key);
        }

        let response = request
            .send()
            .await
            .map_err(|e| GatewayError::VectorDb(format!("Failed to batch store vectors: {}", e)))?;

        if !response.status().is_success() {
            return Err(GatewayError::VectorDb(format!(
                "Failed to batch store vectors: {}",
                response.status()
            )));
        }

        debug!("Batch stored {} vectors", points.len());
        Ok(())
    }

    /// Count vectors in collection
    pub async fn count(&self) -> Result<u64> {
        let url = format!("{}/collections/{}", self.url, self.collection);
        let mut request = self.client.get(&url);

        if let Some(api_key) = &self.api_key {
            request = request.header("api-key", api_key);
        }

        let response = request
            .send()
            .await
            .map_err(|e| GatewayError::VectorDb(format!("Failed to get collection info: {}", e)))?;

        if !response.status().is_success() {
            return Err(GatewayError::VectorDb(format!(
                "Failed to get collection info: {}",
                response.status()
            )));
        }

        let result: serde_json::Value = response.json().await.map_err(|e| {
            GatewayError::VectorDb(format!("Failed to parse collection info: {}", e))
        })?;

        if let Some(count) = result["result"]["points_count"].as_u64() {
            Ok(count)
        } else {
            Ok(0)
        }
    }
}

// Placeholder implementations for Weaviate and Pinecone
#[allow(dead_code)]
impl WeaviateStore {
    /// Create new Weaviate store (not implemented)
    pub async fn new(_config: &VectorDbConfig) -> Result<Self> {
        Err(GatewayError::VectorDb(
            "Weaviate not implemented yet".to_string(),
        ))
    }

    /// Store vector (not implemented)
    pub async fn store(
        &self,
        _id: &str,
        _vector: &[f32],
        _metadata: Option<serde_json::Value>,
    ) -> Result<()> {
        Err(GatewayError::VectorDb(
            "Weaviate not implemented yet".to_string(),
        ))
    }

    /// Search vectors (not implemented)
    pub async fn search(
        &self,
        _query_vector: &[f32],
        _limit: usize,
        _threshold: Option<f32>,
    ) -> Result<Vec<SearchResult>> {
        Err(GatewayError::VectorDb(
            "Weaviate not implemented yet".to_string(),
        ))
    }

    /// Delete vector (not implemented)
    pub async fn delete(&self, _id: &str) -> Result<()> {
        Err(GatewayError::VectorDb(
            "Weaviate not implemented yet".to_string(),
        ))
    }

    /// Get vector by ID (not implemented)
    pub async fn get(&self, _id: &str) -> Result<Option<VectorPoint>> {
        Err(GatewayError::VectorDb(
            "Weaviate not implemented yet".to_string(),
        ))
    }

    /// Health check (not implemented)
    pub async fn health_check(&self) -> Result<()> {
        Err(GatewayError::VectorDb(
            "Weaviate not implemented yet".to_string(),
        ))
    }

    /// Close connection
    pub async fn close(&self) -> Result<()> {
        Ok(())
    }

    /// Batch store vectors (not implemented)
    pub async fn batch_store(&self, _points: &[VectorPoint]) -> Result<()> {
        Err(GatewayError::VectorDb(
            "Weaviate not implemented yet".to_string(),
        ))
    }

    /// Count vectors (not implemented)
    pub async fn count(&self) -> Result<u64> {
        Err(GatewayError::VectorDb(
            "Weaviate not implemented yet".to_string(),
        ))
    }
}

#[allow(dead_code)]
impl PineconeStore {
    /// Create new Pinecone store (not implemented)
    pub async fn new(_config: &VectorDbConfig) -> Result<Self> {
        Err(GatewayError::VectorDb(
            "Pinecone not implemented yet".to_string(),
        ))
    }

    /// Store vector (not implemented)
    pub async fn store(
        &self,
        _id: &str,
        _vector: &[f32],
        _metadata: Option<serde_json::Value>,
    ) -> Result<()> {
        Err(GatewayError::VectorDb(
            "Pinecone not implemented yet".to_string(),
        ))
    }

    /// Search vectors (not implemented)
    pub async fn search(
        &self,
        _query_vector: &[f32],
        _limit: usize,
        _threshold: Option<f32>,
    ) -> Result<Vec<SearchResult>> {
        Err(GatewayError::VectorDb(
            "Pinecone not implemented yet".to_string(),
        ))
    }

    /// Delete vector (not implemented)
    pub async fn delete(&self, _id: &str) -> Result<()> {
        Err(GatewayError::VectorDb(
            "Pinecone not implemented yet".to_string(),
        ))
    }

    /// Get vector by ID (not implemented)
    pub async fn get(&self, _id: &str) -> Result<Option<VectorPoint>> {
        Err(GatewayError::VectorDb(
            "Pinecone not implemented yet".to_string(),
        ))
    }

    /// Health check (not implemented)
    pub async fn health_check(&self) -> Result<()> {
        Err(GatewayError::VectorDb(
            "Pinecone not implemented yet".to_string(),
        ))
    }

    /// Close connection
    pub async fn close(&self) -> Result<()> {
        Ok(())
    }

    /// Batch store vectors (not implemented)
    pub async fn batch_store(&self, _points: &[VectorPoint]) -> Result<()> {
        Err(GatewayError::VectorDb(
            "Pinecone not implemented yet".to_string(),
        ))
    }

    /// Count vectors (not implemented)
    pub async fn count(&self) -> Result<u64> {
        Err(GatewayError::VectorDb(
            "Pinecone not implemented yet".to_string(),
        ))
    }
}

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

    #[test]
    fn test_vector_point_creation() {
        let point = VectorPoint {
            id: "test-1".to_string(),
            vector: vec![0.1, 0.2, 0.3],
            metadata: Some(serde_json::json!({"text": "test document"})),
        };

        assert_eq!(point.id, "test-1");
        assert_eq!(point.vector.len(), 3);
        assert!(point.metadata.is_some());
    }

    #[test]
    fn test_search_result_creation() {
        let result = SearchResult {
            id: "result-1".to_string(),
            score: 0.95,
            metadata: Some(serde_json::json!({"category": "test"})),
            vector: None,
        };

        assert_eq!(result.id, "result-1");
        assert_eq!(result.score, 0.95);
        assert!(result.metadata.is_some());
        assert!(result.vector.is_none());
    }
}