tcvectordb 0.1.9

use tcvectordb::{
    VectorDBClient, Document, Index, VectorIndex, FilterIndex, Filter,
    enums::{IndexType, MetricType, FieldType, ReadConsistency},
    index::HNSWParams,
};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 创建客户端
    let client = VectorDBClient::new(
        "http://127.0.0.1:8100",  // VectorDB服务地址
        "root",                   // 用户名
        "your-api-key",          // API密钥
        ReadConsistency::EventualConsistency,
        30,                      // 超时时间（秒）
    )?;

    // 创建数据库
    let database_name = "book";
    let db = client.create_database_if_not_exists(database_name).await?;
    println!("Database '{}' created or already exists", database_name);

    // 创建索引
    let mut index = Index::new();
    
    // 添加向量索引
    let vector_index = VectorIndex::new(
        "vector",
        3,  // 向量维度
        IndexType::HNSW,
        MetricType::COSINE,
        Some(tcvectordb::index::IndexParams::HNSW(HNSWParams::new(16, 200))),
    );
    index.add_vector_index(vector_index)?;

    // 添加主键索引
    let id_index = FilterIndex::new("id", FieldType::String, IndexType::PRIMARY_KEY);
    index.add_filter_index(id_index)?;

    // 添加过滤索引
    let book_name_index = FilterIndex::new("bookName", FieldType::String, IndexType::FILTER);
    index.add_filter_index(book_name_index)?;

    let page_index = FilterIndex::new("page", FieldType::Uint64, IndexType::FILTER);
    index.add_filter_index(page_index)?;

    // 创建集合
    let collection_name = "book_segments";
    let collection = db.create_collection_if_not_exists(
        collection_name,
        3,                           // 分片数
        2,                           // 副本数
        Some("测试集合".to_string()),   // 描述
        Some(index),                 // 索引
        None,                        // 嵌入模型
        None,                        // TTL配置
    ).await?;
    println!("Collection '{}' created or already exists", collection_name);

    // 准备文档数据
    let documents = vec![
        Document::new()
            .with_id("0001")
            .with_vector(vec![0.2123, 0.21, 0.213])
            .with_field("bookName", "西游记")
            .with_field("author", "吴承恩")
            .with_field("page", 21)
            .with_field("segment", "富贵功名，前缘分定，为人切莫欺心。"),
        
        Document::new()
            .with_id("0002")
            .with_vector(vec![0.2123, 0.22, 0.213])
            .with_field("bookName", "西游记")
            .with_field("author", "吴承恩")
            .with_field("page", 22)
            .with_field("segment", "正大光明，忠良善果弥深。些些狂妄天加谴，眼前不遇待时临。"),
        
        Document::new()
            .with_id("0003")
            .with_vector(vec![0.2123, 0.23, 0.213])
            .with_field("bookName", "三国演义")
            .with_field("author", "罗贯中")
            .with_field("page", 23)
            .with_field("segment", "细作探知这个消息，飞报吕布。"),
        
        Document::new()
            .with_id("0004")
            .with_vector(vec![0.2123, 0.24, 0.213])
            .with_field("bookName", "三国演义")
            .with_field("author", "罗贯中")
            .with_field("page", 24)
            .with_field("segment", "布大惊，与陈宫商议。"),
        
        Document::new()
            .with_id("0005")
            .with_vector(vec![0.2123, 0.25, 0.213])
            .with_field("bookName", "三国演义")
            .with_field("author", "罗贯中")
            .with_field("page", 25)
            .with_field("segment", "玄德曰：布乃当今英勇之士，可出迎之。"),
    ];

    // 插入文档
    let upsert_result = collection.upsert(documents, None, true).await?;
    println!("Documents upserted: {:?}", upsert_result);

    // 等待索引构建
    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;

    // 查询文档
    let document_ids = vec!["0001".to_string(), "0002".to_string(), "0003".to_string()];
    let filter = Filter::new("bookName=\"三国演义\"");
    let output_fields = vec!["id".to_string(), "bookName".to_string()];
    
    let query_results = collection.query(
        Some(document_ids),
        true,                    // 返回向量
        Some(2),                 // 限制数量
        Some(1),                 // 偏移量
        Some(filter),            // 过滤条件
        Some(output_fields),     // 输出字段
        None,                    // 排序
    ).await?;
    
    println!("Query results:");
    for doc in &query_results {
        println!("  ID: {:?}, BookName: {:?}", doc.get_id(), doc.get("bookName"));
    }

    // 向量搜索
    let search_vectors = vec![vec![0.3123, 0.43, 0.213], vec![0.233, 0.12, 0.97]];
    let search_params = tcvectordb::document::SearchParams::new().with_ef(200);
    
    let search_results = collection.search(
        search_vectors,
        None,                    // 无过滤条件
        Some(search_params),     // 搜索参数
        false,                   // 不返回向量
        10,                      // Top K
        None,                    // 输出所有字段
        None,                    // 无超时
        None,                    // 无半径限制
    ).await?;
    
    println!("Search results:");
    for (i, batch) in search_results.iter().enumerate() {
        println!("  Vector {} results:", i + 1);
        for doc in batch {
            println!("    ID: {:?}, BookName: {:?}, Score: {:?}", 
                    doc.get_id(), doc.get("bookName"), doc.get_score());
        }
    }

    // 根据ID搜索
    let search_by_id_results = collection.search_by_id(
        vec!["0003".to_string()],
        Some(Filter::new("bookName=\"三国演义\"")),
        Some(tcvectordb::document::SearchParams::new().with_ef(100)),
        false,
        2,
        None,
        None,
        None,
    ).await?;
    
    println!("Search by ID results:");
    for batch in &search_by_id_results {
        for doc in batch {
            println!("  ID: {:?}, BookName: {:?}", doc.get_id(), doc.get("bookName"));
        }
    }

    // 更新文档
    let update_doc = Document::new().with_field("page", 24);
    let update_result = collection.update(
        update_doc,
        Some(vec!["0001".to_string(), "0003".to_string()]),
        Some(Filter::new("bookName=\"三国演义\"")),
    ).await?;
    println!("Update result: {:?}", update_result);

    // 删除文档
    let delete_result = collection.delete(
        Some(vec!["0001".to_string()]),
        Some(Filter::new("bookName=\"西游记\"")),
        None,
    ).await?;
    println!("Delete result: {:?}", delete_result);

    // 统计文档数量
    let count = collection.count(None).await?;
    println!("Total documents: {}", count);

    // 重建索引
    let rebuild_result = collection.rebuild_index().await?;
    println!("Rebuild index result: {:?}", rebuild_result);

    // 清理：删除集合和数据库
    let drop_collection_result = db.drop_collection(collection_name).await?;
    println!("Collection dropped: {:?}", drop_collection_result);

    let drop_database_result = client.drop_database(database_name).await?;
    println!("Database dropped: {:?}", drop_database_result);

    Ok(())
}