using Api.Models.Rag;

namespace Api.Services.Rag;

public interface ICvVectorStore
{
    void Save(string documentId, IEnumerable<StoredCvChunk> chunks);
    IReadOnlyList<StoredCvChunk> Get(string documentId);
    IReadOnlyList<RetrievedCvChunk> Search(string documentId, float[] queryEmbedding, int topK);
}

public sealed class InMemoryCvVectorStore : ICvVectorStore
{
    private readonly object _lock = new();
    private readonly Dictionary<string, List<StoredCvChunk>> _store = new(StringComparer.OrdinalIgnoreCase);

    public void Save(string documentId, IEnumerable<StoredCvChunk> chunks)
    {
        lock (_lock)
        {
            CleanupExpiredUnsafe();
            _store[documentId] = chunks.ToList();
        }
    }

    public IReadOnlyList<StoredCvChunk> Get(string documentId)
    {
        lock (_lock)
        {
            CleanupExpiredUnsafe();
            return _store.TryGetValue(documentId, out var chunks) ? chunks.ToList() : [];
        }
    }

    public IReadOnlyList<RetrievedCvChunk> Search(string documentId, float[] queryEmbedding, int topK)
    {
        var chunks = Get(documentId);
        if (chunks.Count == 0) return [];

        return chunks
            .Select(chunk => new RetrievedCvChunk
            {
                Text = chunk.Text,
                ChunkIndex = chunk.ChunkIndex,
                Score = CosineSimilarity(queryEmbedding, chunk.Embedding)
            })
            .OrderByDescending(x => x.Score)
            .Take(Math.Clamp(topK, 1, 12))
            .ToList();
    }

    private void CleanupExpiredUnsafe()
    {
        var now = DateTimeOffset.UtcNow;
        foreach (var key in _store.Where(x => x.Value.All(c => c.ExpiresAt <= now)).Select(x => x.Key).ToList())
        {
            _store.Remove(key);
        }
    }

    private static double CosineSimilarity(float[] a, float[] b)
    {
        if (a.Length != b.Length || a.Length == 0) return 0;

        double dot = 0;
        double magA = 0;
        double magB = 0;

        for (var i = 0; i < a.Length; i++)
        {
            dot += a[i] * b[i];
            magA += a[i] * a[i];
            magB += b[i] * b[i];
        }

        if (magA == 0 || magB == 0) return 0;
        return dot / (Math.Sqrt(magA) * Math.Sqrt(magB));
    }
}