Nexus.Reader/src/NexusReader.Infrastructure/Persistence/VectorSearchStore.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.Logging;
using Qdrant.Client;
using Qdrant.Client.Grpc;
using Polly;
using Polly.Registry;
using NexusReader.Application.Abstractions.Persistence;

namespace NexusReader.Infrastructure.Persistence;

/// <summary>
/// Infrastructure implementation of <see cref="IVectorSearchStore"/> utilizing <see cref="QdrantClient"/> 
/// and <see cref="IEmbeddingGenerator{TInput, TEmbedding}"/> to execute semantic vector queries.
/// </summary>
internal sealed class VectorSearchStore : IVectorSearchStore
{
    private readonly QdrantClient _qdrantClient;
    private readonly IEmbeddingGenerator<string, Embedding<float>> _embeddingGenerator;
    private readonly ResiliencePipeline _retryPipeline;
    private readonly ILogger<VectorSearchStore> _logger;

    public VectorSearchStore(
        QdrantClient qdrantClient,
        IEmbeddingGenerator<string, Embedding<float>> embeddingGenerator,
        ResiliencePipelineProvider<string> pipelineProvider,
        ILogger<VectorSearchStore> logger)
    {
        _qdrantClient = qdrantClient;
        _embeddingGenerator = embeddingGenerator;
        _retryPipeline = pipelineProvider.GetPipeline("ai-retry");
        _logger = logger;
    }

    /// <inheritdoc />
    public async Task<List<VectorChunk>> SearchGlobalAsync(string queryText, string tenantId, int limit, CancellationToken cancellationToken = default)
    {
        var queryVector = await GenerateEmbeddingAsync(queryText, cancellationToken);
        var filter = BuildTenantFilter(tenantId);

        return await ExecuteSearchAsync(queryVector, filter, limit, cancellationToken);
    }

    /// <inheritdoc />
    public async Task<List<VectorChunk>> SearchLocalAsync(string queryText, string tenantId, List<Guid> whitelistedBookIds, int limit, CancellationToken cancellationToken = default)
    {
        if (whitelistedBookIds == null || !whitelistedBookIds.Any())
        {
            return new List<VectorChunk>();
        }

        var queryVector = await GenerateEmbeddingAsync(queryText, cancellationToken);
        var filter = BuildTenantFilter(tenantId);

        var whitelistFilter = new Qdrant.Client.Grpc.Filter();
        foreach (var bookId in whitelistedBookIds)
        {
            whitelistFilter.Should.Add(new Qdrant.Client.Grpc.Condition
            {
                Field = new Qdrant.Client.Grpc.FieldCondition
                {
                    Key = "ebookId",
                    Match = new Qdrant.Client.Grpc.Match { Text = bookId.ToString() }
                }
            });
        }
        filter.Must.Add(new Qdrant.Client.Grpc.Condition { Filter = whitelistFilter });

        return await ExecuteSearchAsync(queryVector, filter, limit, cancellationToken);
    }

    /// <inheritdoc />
    public async Task<List<VectorChunk>> SearchGlobalExcludeAsync(string queryText, string tenantId, Guid excludeBookId, int limit, CancellationToken cancellationToken = default)
    {
        var queryVector = await GenerateEmbeddingAsync(queryText, cancellationToken);
        var filter = BuildTenantFilter(tenantId);

        // Exclude current book
        filter.MustNot.Add(new Qdrant.Client.Grpc.Condition
        {
            Field = new Qdrant.Client.Grpc.FieldCondition
            {
                Key = "ebookId",
                Match = new Qdrant.Client.Grpc.Match { Text = excludeBookId.ToString() }
            }
        });

        return await ExecuteSearchAsync(queryVector, filter, limit, cancellationToken);
    }

    private async Task<float[]> GenerateEmbeddingAsync(string text, CancellationToken cancellationToken)
    {
        if (string.IsNullOrWhiteSpace(text))
        {
            _logger.LogWarning("[VectorSearchStore] Attempted to generate embedding from empty text. Returning zero vector.");
            return Array.Empty<float>();
        }

        var sw = Stopwatch.StartNew();
        var response = await _retryPipeline.ExecuteAsync(async ct => 
            await _embeddingGenerator.GenerateAsync(
                new[] { text }, 
                new EmbeddingGenerationOptions { Dimensions = 768 }, 
                cancellationToken: ct), cancellationToken);
        sw.Stop();

        _logger.LogDebug("[VectorSearchStore] Embedding generated in {ElapsedMs}ms for text of {Length} chars.", sw.ElapsedMilliseconds, text.Length);
        return response.First().Vector.ToArray();
    }

    private Qdrant.Client.Grpc.Filter BuildTenantFilter(string tenantId)
    {
        var filter = new Qdrant.Client.Grpc.Filter();
        var tenantFilter = new Qdrant.Client.Grpc.Filter();

        tenantFilter.Should.Add(new Qdrant.Client.Grpc.Condition
        {
            Field = new Qdrant.Client.Grpc.FieldCondition
            {
                Key = "tenantId",
                Match = new Qdrant.Client.Grpc.Match { Text = tenantId }
            }
        });

        tenantFilter.Should.Add(new Qdrant.Client.Grpc.Condition
        {
            Field = new Qdrant.Client.Grpc.FieldCondition
            {
                Key = "tenantId",
                Match = new Qdrant.Client.Grpc.Match { Text = "global" }
            }
        });

        filter.Must.Add(new Qdrant.Client.Grpc.Condition { Filter = tenantFilter });
        return filter;
    }

    private async Task<List<VectorChunk>> ExecuteSearchAsync(float[] queryVector, Qdrant.Client.Grpc.Filter filter, int limit, CancellationToken cancellationToken)
    {
        if (queryVector.Length == 0)
        {
            _logger.LogWarning("[VectorSearchStore] Empty query vector — skipping Qdrant search.");
            return new List<VectorChunk>();
        }

        try
        {
            await EnsureCollectionExistsAsync("knowledge_units", cancellationToken);

            var sw = Stopwatch.StartNew();
            var response = await _qdrantClient.SearchAsync(
                collectionName: "knowledge_units",
                vector: queryVector,
                filter: filter,
                limit: (ulong)limit,
                cancellationToken: cancellationToken
            );
            sw.Stop();
            _logger.LogInformation("[VectorSearchStore] Qdrant search returned {Count} results in {ElapsedMs}ms.", response.Count, sw.ElapsedMilliseconds);

            return response.Select(point => 
            {
                var content = point.Payload.TryGetValue("content", out var cv) ? cv.StringValue : string.Empty;
                var ebookId = point.Payload.TryGetValue("ebookId", out var ev) ? ev.StringValue : string.Empty;
                var metadataJson = point.Payload.TryGetValue("metadataJson", out var mv) ? mv.StringValue : string.Empty;
                var bookTitle = point.Payload.TryGetValue("bookTitle", out var btv) ? btv.StringValue : string.Empty;
                var chapterTitle = point.Payload.TryGetValue("chapterTitle", out var ctv) ? ctv.StringValue : string.Empty;

                return new VectorChunk(content, ebookId, point.Score, metadataJson, bookTitle, chapterTitle);
            }).ToList();
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "[VectorSearchStore] Qdrant search execution failed.");
            throw;
        }
    }

    private async Task EnsureCollectionExistsAsync(string collectionName, CancellationToken cancellationToken)
    {
        try
        {
            var exists = await _qdrantClient.CollectionExistsAsync(collectionName, cancellationToken);
            if (!exists)
            {
                _logger.LogInformation("[VectorSearchStore] Collection '{CollectionName}' does not exist — creating.", collectionName);
                await _qdrantClient.CreateCollectionAsync(
                    collectionName: collectionName,
                    vectorsConfig: new Qdrant.Client.Grpc.VectorParams
                    {
                        Size = 768,
                        Distance = Distance.Cosine
                    },
                    cancellationToken: cancellationToken
                );
                _logger.LogInformation("[VectorSearchStore] Collection '{CollectionName}' created successfully.", collectionName);
            }
        }
        catch (Exception ex)
        {
            // Log concurrent creation conflicts (e.g., AlreadyExists gRPC status) but do not propagate.
            _logger.LogWarning(ex, "[VectorSearchStore] Non-fatal error while ensuring collection '{CollectionName}' exists. Possible concurrent creation.", collectionName);
        }
    }
}