hwhkit 0.1.2

//! Qdrant 向量数据库连接管理模块

use crate::config::QdrantConfig;
use crate::error::{Error, Result};
use qdrant_client::client::QdrantClient;
use qdrant_client::qdrant::HealthCheckReply;
use std::sync::Arc;
use std::time::Duration;

/// 向量数据库管理器
#[derive(Clone)]
pub struct VectorDatabase {
    client: Arc<QdrantClient>,
    config: QdrantConfig,
}

impl std::fmt::Debug for VectorDatabase {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("VectorDatabase")
            .field("config", &self.config)
            .finish()
    }
}

impl VectorDatabase {
    /// 从配置创建向量数据库连接
    ///
    /// # Arguments
    ///
    /// * `config` - 向量数据库配置
    ///
    /// # Returns
    ///
    /// 返回向量数据库实例或错误
    pub async fn from_config(config: QdrantConfig) -> Result<Self> {
        if !config.enabled {
            return Err(Error::Database("Qdrant 向量数据库未启用".to_string()));
        }

        // 创建客户端配置
        let mut client_builder = QdrantClient::from_url(&config.url)
            .with_timeout(Duration::from_secs(config.timeout));

        // 如果提供了 API Key，设置认证
        if let Some(api_key) = &config.api_key {
            if !api_key.is_empty() {
                client_builder = client_builder.with_api_key(api_key.clone());
            }
        }

        let client = client_builder
            .build()
            .map_err(|e| Error::Database(format!("Qdrant 连接失败: {}", e)))?;

        tracing::info!(
            "Qdrant 向量数据库连接成功: {}, 默认集合: {}",
            config.url,
            config.default_collection
        );

        Ok(Self {
            client: Arc::new(client),
            config,
        })
    }

    /// 获取 Qdrant 客户端
    pub fn client(&self) -> &QdrantClient {
        &self.client
    }

    /// 获取配置
    pub fn config(&self) -> &QdrantConfig {
        &self.config
    }

    /// 获取默认集合名称
    pub fn default_collection(&self) -> &str {
        &self.config.default_collection
    }

    /// 测试数据库连接
    pub async fn ping(&self) -> Result<HealthCheckReply> {
        self.client
            .health_check()
            .await
            .map_err(|e| Error::Database(format!("Qdrant 连接测试失败: {}", e)))
    }

    /// 检查集合是否存在
    pub async fn collection_exists(&self, collection_name: &str) -> Result<bool> {
        self.client
            .collection_exists(collection_name)
            .await
            .map_err(|e| Error::Database(format!("检查集合失败: {}", e)))
    }

    /// 列出所有集合
    pub async fn list_collections(&self) -> Result<Vec<String>> {
        let response = self
            .client
            .list_collections()
            .await
            .map_err(|e| Error::Database(format!("列出集合失败: {}", e)))?;

        Ok(response
            .collections
            .iter()
            .map(|c| c.name.clone())
            .collect())
    }
}

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

    #[tokio::test]
    #[ignore] // 需要实际的 Qdrant 服务
    async fn test_vector_database_connection() {
        let config = QdrantConfig {
            enabled: true,
            url: "http://localhost:6334".to_string(),
            api_key: None,
            timeout: 30,
            default_collection: "test".to_string(),
        };

        let db = VectorDatabase::from_config(config).await;
        assert!(db.is_ok());

        if let Ok(db) = db {
            assert!(db.ping().await.is_ok());
        }
    }
}