using Microsoft.Extensions.AI;
using Microsoft.Extensions.VectorData;
using Microsoft.SemanticKernel.Connectors.Sqlite;

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

List<Dog> dogs =
[
    new()
    {
        Key = 1,
        Name = "Dorka",
        Behavior = "Dorka is perfect example of Dachshund breed. She's smart, curious and brave. Smart, but sometimes stubborn. Likes people and is a great dog for kids and as a therapy dog. Dorka enjoys digging and burrowing. She can be quite vocal when she's excited or wants attention."
    },
    new()
    {
        Key = 2,
        Name = "Max",
        Behavior = "Max is energetic and playful. He loves to run around and play fetch. He's very friendly and gets along well with other dogs and people. Max is very social and enjoys going to dog parks. He's curious and loves exploring new places."
    },
    new()
    {
        Key = 3,
        Name = "Bella",
        Behavior = "Bella is patient and responds well to gentle training methods. She's very obedient and eager to please. Bella loves to lounge around the house and enjoys a cozy spot to nap. She's not as high-energy as Max but still enjoys her daily walks."
    },
];
var vectorStore = new SqliteVectorStore("Data Source=demo.db");
var dogsCollection = vectorStore.GetCollection<string, Dog>("dog");
await dogsCollection.CreateCollectionIfNotExistsAsync();
foreach (var dog in dogs)
{
    dog.Embedding = await generator.GenerateVectorAsync(dog.Behavior);
    await dogsCollection.UpsertAsync(dog);
}

var query = "A family friendly dog.";
var queryEmbedding = await generator.GenerateVectorAsync(query);
var searchOptions = new VectorSearchOptions<Dog>()
{
    //Filter = x => x.Key < 100,
};
Console.WriteLine("Searching...");
Console.WriteLine();
await foreach (var result in dogsCollection.SearchEmbeddingAsync(queryEmbedding, 2, searchOptions))
{
    Console.WriteLine($"Name: {result.Record.Name}");
    Console.WriteLine($"Behavior: {result.Record.Behavior}");
    Console.WriteLine($"Score: {result.Score}");
    Console.WriteLine();
}

public class Dog
{
    [VectorStoreRecordKey]
    public ulong Key { get; set; }

    [VectorStoreRecordData]
    public string Name { get; set; } = null!;

    [VectorStoreRecordData]
    public string Behavior { get; set; } = null!;

    [VectorStoreRecordVector(384, DistanceFunction = DistanceFunction.CosineDistance)]
    public ReadOnlyMemory<float> Embedding { get; set; }
}
