// Copyright (c) Microsoft. All rights reserved.

// This sample shows how to use Qdrant with a custom schema to add retrieval augmented generation (RAG) capabilities to an AI agent.
// While the sample is using Qdrant, it can easily be replaced with any other vector store that implements the Microsoft.Extensions.VectorData abstractions.
// The TextSearchProvider runs a search against the vector store before each model invocation and injects the results into the model context.

using Azure.AI.Projects;
using Azure.Identity;
using Microsoft.Agents.AI;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.VectorData;
using Microsoft.SemanticKernel.Connectors.Qdrant;
using Qdrant.Client;

var endpoint = Environment.GetEnvironmentVariable("FOUNDRY_PROJECT_ENDPOINT") ?? throw new InvalidOperationException("FOUNDRY_PROJECT_ENDPOINT is not set.");
var deploymentName = Environment.GetEnvironmentVariable("FOUNDRY_MODEL") ?? "gpt-5.4-mini";
var embeddingDeploymentName = Environment.GetEnvironmentVariable("FOUNDRY_EMBEDDING_MODEL") ?? "text-embedding-3-large";
var afOverviewUrl = "https://raw.githubusercontent.com/MicrosoftDocs/semantic-kernel-docs/refs/heads/main/agent-framework/overview/index.md";
var afMigrationUrl = "https://raw.githubusercontent.com/MicrosoftDocs/semantic-kernel-docs/refs/heads/main/agent-framework/migration-guide/from-semantic-kernel/index.md";

// WARNING: DefaultAzureCredential is convenient for development but requires careful consideration in production.
// In production, consider using a specific credential (e.g., ManagedIdentityCredential) to avoid
// latency issues, unintended credential probing, and potential security risks from fallback mechanisms.
AIProjectClient aiProjectClient = new(
    new Uri(endpoint),
    new DefaultAzureCredential());

// Create a Qdrant vector store that uses the Azure AI Foundry embedding model to generate embeddings.
QdrantClient client = new("localhost");
VectorStore vectorStore = new QdrantVectorStore(client, ownsClient: true, new()
{
    EmbeddingGenerator = aiProjectClient.GetProjectOpenAIClient().GetEmbeddingClient(embeddingDeploymentName).AsIEmbeddingGenerator()
});

// Create a collection and upsert some text into it.
var documentationCollection = vectorStore.GetCollection<Guid, DocumentationChunk>("documentation");
await documentationCollection.EnsureCollectionDeletedAsync(); // Clear out any data from previous runs.
await documentationCollection.EnsureCollectionExistsAsync();
await UploadDataFromMarkdown(afOverviewUrl, "Microsoft Agent Framework Overview", documentationCollection, 2000, 200);
await UploadDataFromMarkdown(afMigrationUrl, "Semantic Kernel to Microsoft Agent Framework Migration Guide", documentationCollection, 2000, 200);

// Create an adapter function that the TextSearchProvider can use to run searches against the collection.
Func<string, CancellationToken, Task<IEnumerable<TextSearchProvider.TextSearchResult>>> SearchAdapter = async (text, ct) =>
{
    List<TextSearchProvider.TextSearchResult> results = [];
    await foreach (var result in documentationCollection.SearchAsync(text, 5, cancellationToken: ct))
    {
        results.Add(new TextSearchProvider.TextSearchResult
        {
            SourceName = result.Record.SourceName,
            SourceLink = result.Record.SourceLink,
            Text = result.Record.Text ?? string.Empty,
            RawRepresentation = result
        });
    }
    return results;
};

// Configure the options for the TextSearchProvider.
TextSearchProviderOptions textSearchOptions = new()
{
    // Run the search prior to every model invocation.
    SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke,
    // Use up to 5 recent messages when searching so that searches
    // still produce valuable results even when the user is referring
    // back to previous messages in their request.
    RecentMessageMemoryLimit = 5
};

// Create the AI agent with the TextSearchProvider as the AI context provider.
AIAgent agent = aiProjectClient
    .AsAIAgent(new ChatClientAgentOptions
    {
        ChatOptions = new() { ModelId = deploymentName, Instructions = "You are a helpful support specialist for the Microsoft Agent Framework. Answer questions using the provided context and cite the source document when available. Keep responses brief." },
        AIContextProviders = [new TextSearchProvider(SearchAdapter, textSearchOptions)],
        // Configure a filter on the InMemoryChatHistoryProvider so that we don't persist the messages produced by the TextSearchProvider in chat history.
        // The default is to persist all messages except those that came from chat history in the first place.
        // You may choose to persist the TextSearchProvider messages, if you want the search output to be provided to the model in future interactions as well.
        ChatHistoryProvider = new InMemoryChatHistoryProvider(new InMemoryChatHistoryProviderOptions()
        {
            StorageInputRequestMessageFilter = msgs => msgs.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory && m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.AIContextProvider)
        })
    });

AgentSession session = await agent.CreateSessionAsync();

Console.WriteLine(">> Asking about SK sessions\n");
Console.WriteLine(await agent.RunAsync("Hi! How do I create a thread/session in Semantic Kernel?", session));

// Here we are asking a very vague question when taken out of context,
// but since we are including previous messages in our search using RecentMessageMemoryLimit
// the RAG search should still produce useful results.
Console.WriteLine("\n>> Asking about AF sessions\n");
Console.WriteLine(await agent.RunAsync("and in Agent Framework?", session));

Console.WriteLine("\n>> Contrasting Approaches\n");
Console.WriteLine(await agent.RunAsync("Please contrast the two approaches", session));

Console.WriteLine("\n>> Asking about ancestry\n");
Console.WriteLine(await agent.RunAsync("What are the predecessors to the Agent Framework?", session));

static async Task UploadDataFromMarkdown(string markdownUrl, string sourceName, VectorStoreCollection<Guid, DocumentationChunk> vectorStoreCollection, int chunkSize, int overlap)
{
    // Download the markdown from the given url.
    using HttpClient client = new();
    var markdown = await client.GetStringAsync(new Uri(markdownUrl));

    // Chunk it into separate parts with some overlap between chunks
    var chunks = new List<DocumentationChunk>();
    for (int i = 0; i < markdown.Length; i += chunkSize)
    {
        var chunk = new DocumentationChunk
        {
            Key = Guid.NewGuid(),
            SourceLink = markdownUrl,
            SourceName = sourceName,
            Text = markdown.Substring(i, Math.Min(chunkSize + overlap, markdown.Length - i))
        };
        chunks.Add(chunk);
    }

    // Upsert each chunk into the provided vector store.
    await vectorStoreCollection.UpsertAsync(chunks);
}

// Data model that defines the database schema we want to use.
internal sealed class DocumentationChunk
{
    [VectorStoreKey]
    public Guid Key { get; set; }
    [VectorStoreData]
    public string SourceLink { get; set; } = string.Empty;
    [VectorStoreData]
    public string SourceName { get; set; } = string.Empty;
    [VectorStoreData]
    public string Text { get; set; } = string.Empty;
    [VectorStoreVector(Dimensions: 3072)]
    public string Embedding => this.Text;
}
