using ChromaDB.Client;
using Microsoft.Extensions.AI;

var options = new ChromaConfigurationOptions(uri: "http://192.168.66.34:8000/api/v1/");
using var httpClient = new HttpClient();
var client = new ChromaClient(options, httpClient);

var collection = await client.GetOrCreateCollection("data");
var collectionClient = new ChromaCollectionClient(collection, options, httpClient);

using var generator = new OllamaEmbeddingGenerator(new Uri("http://192.168.66.34:11434"), "all-minilm");

var ids = new List<string>();
var embeddings = new List<ReadOnlyMemory<float>>();
await foreach (var (id, vector) in GetEmbeddings([
    "Dog is human's best friend.",
    "Cat is smaller than dog.",
    "Tiger is strong.",
    "Cow is bigger than tiger.",
    "Bird can fly.",
    "Elephant is bigger than cow."]))
{
    ids.Add(Guid.NewGuid().ToString());
    embeddings.Add(vector);
}
await collectionClient.Add(ids, embeddings);

var mouseVector = await generator.GenerateVectorAsync("Big Mouse");
var results = await collectionClient.Query(mouseVector,
    include: ChromaQueryInclude.Distances | ChromaQueryInclude.Metadatas,
    nResults: 3,
    where: null);
foreach (var item in results)
{
    Console.WriteLine($"[{item.Id}]: D={item.Distance}");
}

async IAsyncEnumerable<(string, ReadOnlyMemory<float>)> GetEmbeddings(IEnumerable<string> items)
{
    foreach (var item in items)
    {
        var vector = await generator.GenerateVectorAsync(item);
        yield return (item, vector);
    }
}