tcvectordb 0.1.9

use tcvectordb::{
    enums::{FieldType, IndexType, MetricType, ReadConsistency},
    index::HNSWParams,
    Embedding, FilterIndex, Index, Result, VectorDBClient, VectorIndex,
};
use uuid;

#[tokio::main]
async fn main() -> Result<()> {
    println!("📚 Python Equivalent Demo - Creating Collection with Embedding Config");

    // 设置连接参数
    let url = std::env::var("VECTORDB_URL").unwrap_or_else(|_| "http://localhost:8100".to_string());
    let username = std::env::var("VECTORDB_USERNAME").unwrap_or_else(|_| "root".to_string());
    let api_key = std::env::var("VECTORDB_API_KEY")
        .expect("VECTORDB_API_KEY environment variable is required");

    // 创建客户端 - 对应 Python: client = tcvectordb.VectorDBClient(...)
    let client = VectorDBClient::new(
        &url,
        &username,
        &api_key,
        ReadConsistency::EventualConsistency,
        30,
    )?;

    println!("✅ VectorDB client created successfully");

    // 创建数据库 - 对应 Python: client.create_database_if_not_exists('db-test')
    let db_name = "db-test";
    let db = client.create_database_if_not_exists(db_name).await?;
    println!("✅ Database '{}' ready", db_name);

    // 索引配置 - 对应 Python 的 index = Index(...)
    let mut index = Index::new();

    // 对应 Python: FilterIndex(name='id', field_type=FieldType.String, index_type=IndexType.PRIMARY_KEY)
    index.add_filter_index(FilterIndex::new(
        "id",
        FieldType::String,
        IndexType::PRIMARY_KEY,
    ))?;

    // 对应 Python: VectorIndex(name='vector', dimension=768, index_type=IndexType.HNSW, ...)
    index.add_vector_index(VectorIndex::new(
        "vector",
        768, // bge-base-zh模型的维度
        IndexType::HNSW,
        MetricType::COSINE,
        Some(tcvectordb::index::IndexParams::HNSW(HNSWParams::new(
            16, 200,
        ))),
    ))?;

    // 对应 Python: FilterIndex(name='author', field_type=FieldType.String, index_type=IndexType.FILTER)
    index.add_filter_index(FilterIndex::new(
        "author",
        FieldType::String,
        IndexType::FILTER,
    ))?;

    // 对应 Python: FilterIndex(name='tags', field_type=FieldType.Array, index_type=IndexType.FILTER)
    index.add_filter_index(FilterIndex::new(
        "tags",
        FieldType::Array,
        IndexType::FILTER,
    ))?;

    // 对应 Python: FilterIndex(name='bookName', field_type=FieldType.String, index_type=IndexType.FILTER)
    index.add_filter_index(FilterIndex::new(
        "bookName",
        FieldType::String,
        IndexType::FILTER,
    ))?;

    // Embedding 配置 - 对应 Python: ebd = Embedding(vector_field='vector', field='text', model_name='bge-base-zh')
    let embedding = Embedding::new("vector", "text").with_model_name("bge-base-zh");

    println!("🔧 Creating collection with embedding configuration...");

    // 使用唯一的集合名称
    let collection_name = format!(
        "book-emb-{}",
        uuid::Uuid::new_v4().to_string().replace("-", "")[..8].to_string()
    );

    // 创建集合 - 对应 Python: coll = client.create_collection(...)
    let collection = db
        .create_collection(
            &collection_name, // collection_name='book-emb'
            1,                // shard=1
            0,                // replicas=1 (但服务器要求为0)
            Some("this is a collection of test embedding".to_string()), // description
            Some(index),      // index=index
            Some(embedding),  // embedding=ebd
            None,             // ttl_config
        )
        .await?;

    println!(
        "✅ Collection '{}' created successfully!",
        collection.name()
    );

    // 打印集合信息 - 对应 Python: print(vars(coll))
    println!("📋 Collection Information:");
    println!("   - Name: {}", collection.name());
    println!("   - Database: {}", collection.database().name());
    println!(
        "   - Embedding configured: ✅ (vector_field='vector', field='text', model='bge-base-zh')"
    );
    println!("   - Index configured: ✅ (768-dim HNSW vector + filter indexes)");

    // 注意：由于配置了嵌入模型，这个集合可以：
    // 1. 自动将text字段转换为768维向量存储在vector字段
    // 2. 支持文本搜索（search_by_text）
    // 3. 支持混合搜索

    println!("\n💡 This collection is now configured with:");
    println!("   - Automatic text-to-vector embedding using bge-base-zh model");
    println!("   - 768-dimensional HNSW vector index");
    println!("   - Filter indexes for author, tags, and bookName fields");
    println!("   - Ready for text search and hybrid search operations");

    // 清理
    println!("\n🧹 Cleaning up...");
    match db.drop_collection(&collection_name).await {
        Ok(_) => println!("✅ Collection '{}' dropped", collection_name),
        Err(e) => println!("⚠️  Failed to drop collection: {}", e),
    }

    match client.drop_database(db_name).await {
        Ok(_) => println!("✅ Database '{}' dropped", db_name),
        Err(e) => println!("⚠️  Failed to drop database: {}", e),
    }

    println!("\n🎉 Python equivalent demo completed!");
    println!("📝 This Rust code is equivalent to the Python example you provided.");

    Ok(())
}