// TEMPLATE: VectorSearchTool.cs
// PURPOSE: A typed MAF tool that retrieves relevant chunks from the pgvector knowledge base.
// READS STANDARD: ai-agents-rag-custom-pgvector
// PLACEHOLDERS:
//   {{PROJECT_NAMESPACE}}  — root namespace of the project (e.g., MyProject)
// LAST-VERIFIED: 2026-05-19
//
// SECURITY NOTE — tenant_id:
//   tenant_id is resolved from the injected ITenantContext (or equivalent auth accessor),
//   NEVER from a tool parameter visible to the LLM. A model-facing tenant parameter would
//   allow prompt-injection to cross tenant boundaries.
//
// NPGSQL + PGVECTOR PARAMETER BINDING:
//   This template uses NpgsqlDataSource (raw SQL) for the cosine-distance query.
//   The project must call dataSourceBuilder.UseVector() before building the data source
//   so Npgsql can map the vector type — see infrastructure-neon-pgvector.
//
//   // VERIFY: The exact Npgsql parameter binding for a pgvector type has been written
//   // conservatively below using new Pgvector.Vector(float[]).
//   // Confirmed pattern from neon-pgvector standard (Pgvector package):
//   //   cmd.Parameters.AddWithValue("queryVec", new Pgvector.Vector(embedding.ToArray()));
//   // If the project uses a different Npgsql/Pgvector version, verify the parameter API.
//
// REGISTRATION (in agent setup code):
//   var tool = sp.GetRequiredService<VectorSearchTool>();
//   var agent = modelRegistry.GetChatClient("text-default")
//       .AsAIAgent(
//           instructions: systemPrompt,
//           tools: [AIFunctionFactory.Create(tool.SearchAsync)]);
//
// ALTERNATIVE — EF Core:
//   If the project already wires pgvector through EF Core (UseVector(), HasMethod("hnsw")),
//   replace the NpgsqlDataSource usage with db.Database.SqlQuery<string>() per
//   infrastructure-neon-pgvector. The tool method signature stays identical.

using System.ComponentModel;
using Npgsql;
using Pgvector;

namespace {{PROJECT_NAMESPACE}}.AI;

/// <summary>
/// MAF tool: searches the knowledge base for chunks semantically similar to a query.
/// Registered on an agent via <c>AIFunctionFactory.Create(tool.SearchAsync)</c>.
/// The agent calls this autonomously when it needs grounding context.
/// </summary>
public sealed class VectorSearchTool(
    EmbeddingService embeddingService,
    NpgsqlDataSource dataSource,
    ITenantContext tenantContext)   // Resolved from DI — scoped to the authenticated request.
                                    // Replace with Guid currentTenantId if a simpler accessor is preferred.
{
    private const int DefaultTopK = 5;

    /// <summary>
    /// Searches the knowledge base for context relevant to a question.
    /// </summary>
    /// <param name="query">The natural-language search query.</param>
    /// <param name="ct">Cancellation token.</param>
    /// <returns>
    /// Concatenated text of the top matching chunks, separated by double newlines.
    /// Returns an empty string if no results are found.
    /// </returns>
    [Description("Searches the knowledge base for context relevant to a question")]
    public async Task<string> SearchAsync(
        [Description("the natural-language search query")] string query,
        CancellationToken ct = default)
    {
        // 1. Embed the query at search time.
        ReadOnlyMemory<float> queryVector = await embeddingService.EmbedAsync(query, ct);

        // 2. Resolve tenant from auth context — never from a model-supplied parameter.
        Guid tenantId = tenantContext.CurrentTenantId;

        // 3. Run parameterized cosine-distance query.
        //    ORDER BY embedding <=> @queryVec sorts ascending by distance (most similar first).
        const string sql = """
            SELECT content
            FROM   kb_embeddings
            WHERE  tenant_id = @tenant
            ORDER  BY embedding <=> @queryVec
            LIMIT  @topK
            """;

        var chunks = new List<string>(DefaultTopK);

        await using var conn = await dataSource.OpenConnectionAsync(ct);
        await using var cmd = new NpgsqlCommand(sql, conn);

        cmd.Parameters.AddWithValue("tenant", tenantId);

        // VERIFY: Pgvector.Vector wraps a float[] for Npgsql parameter binding.
        // The project must have called dataSourceBuilder.UseVector() on NpgsqlDataSourceBuilder.
        // NuGet: Pgvector (matches the Pgvector.EntityFrameworkCore used in neon-pgvector.md).
        cmd.Parameters.AddWithValue("queryVec", new Vector(queryVector.ToArray()));

        cmd.Parameters.AddWithValue("topK", DefaultTopK);

        await using var reader = await cmd.ExecuteReaderAsync(ct);
        while (await reader.ReadAsync(ct))
        {
            chunks.Add(reader.GetString(0));
        }

        return string.Join("\n\n", chunks);
    }
}