ChatRAG/Data/ChromaProjectDataRepository.cs

using ChatRAG.Models;
using ChatRAG.Services.Contracts;
using ChatRAG.Services.SearchVectors;
using ChatRAG.Settings.ChatRAG.Configuration;
using Microsoft.Extensions.Options;
using System.Text;
using System.Text.Json;

namespace ChatRAG.Data
{

    public class ChromaProjectDataRepository : IProjectDataRepository
    {
        private readonly HttpClient _httpClient;
        private readonly string _collectionName;
        private readonly ILogger<ChromaProjectDataRepository> _logger;

        public ChromaProjectDataRepository(
            IOptions<VectorDatabaseSettings> settings,
            HttpClient httpClient,
            ILogger<ChromaProjectDataRepository> logger)
        {
            var chromaSettings = settings.Value.Chroma ?? throw new ArgumentNullException("Chroma settings not configured");
            _httpClient = httpClient;
            _httpClient.BaseAddress = new Uri($"http://{chromaSettings.Host}:{chromaSettings.Port}");
            _collectionName = "projects"; // Collection separada para projetos
            _logger = logger;

            InitializeAsync().GetAwaiter().GetResult();
        }

        private async Task InitializeAsync()
        {
            try
            {
                // Verificar se a collection existe, se não, criar
                var collections = await GetCollectionsAsync();
                if (!collections.Contains(_collectionName))
                {
                    await CreateProjectsCollection();
                }
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Erro ao inicializar collection de projetos no Chroma");
            }
        }

        public async Task<List<Project>> GetAsync()
        {
            try
            {
                var query = new
                {
                    where = new { entity_type = "project" },
                    include = new[] { "documents", "metadatas" }
                };

                var json = JsonSerializer.Serialize(query);
                var content = new StringContent(json, Encoding.UTF8, "application/json");

                var response = await _httpClient.PostAsync($"/api/v2/collections/{_collectionName}/get", content);

                if (!response.IsSuccessStatusCode)
                {
                    _logger.LogError("Erro ao buscar projetos no Chroma");
                    return new List<Project>();
                }

                var result = await response.Content.ReadAsStringAsync();
                var getResult = JsonSerializer.Deserialize<ChromaGetResult>(result);

                var projects = new List<Project>();
                if (getResult?.ids?.Length > 0)
                {
                    for (int i = 0; i < getResult.ids.Length; i++)
                    {
                        var metadata = getResult.metadatas?[i];
                        if (metadata != null)
                        {
                            projects.Add(new Project
                            {
                                Id = metadata.GetValueOrDefault("id")?.ToString() ?? getResult.ids[i],
                                Nome = metadata.GetValueOrDefault("nome")?.ToString() ?? "",
                                Descricao = metadata.GetValueOrDefault("descricao")?.ToString() ?? ""
                            });
                        }
                    }
                }

                return projects;
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Erro ao recuperar projetos do Chroma");
                return new List<Project>();
            }
        }

        public async Task<Project?> GetAsync(string id)
        {
            try
            {
                var query = new
                {
                    ids = new[] { id },
                    include = new[] { "metadatas" }
                };

                var json = JsonSerializer.Serialize(query);
                var content = new StringContent(json, Encoding.UTF8, "application/json");

                var response = await _httpClient.PostAsync($"/api/v2/collections/{_collectionName}/get", content);

                if (!response.IsSuccessStatusCode)
                {
                    return null;
                }

                var result = await response.Content.ReadAsStringAsync();
                var getResult = JsonSerializer.Deserialize<ChromaGetResult>(result);

                if (getResult?.ids?.Length > 0 && getResult.metadatas?[0] != null)
                {
                    var metadata = getResult.metadatas[0];
                    return new Project
                    {
                        Id = metadata.GetValueOrDefault("id")?.ToString() ?? id,
                        Nome = metadata.GetValueOrDefault("nome")?.ToString() ?? "",
                        Descricao = metadata.GetValueOrDefault("descricao")?.ToString() ?? ""
                    };
                }

                return null;
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Erro ao buscar projeto {Id} no Chroma", id);
                return null;
            }
        }

        public async Task CreateAsync(Project newProject)
        {
            try
            {
                var id = string.IsNullOrEmpty(newProject.Id) ? Guid.NewGuid().ToString() : newProject.Id;
                newProject.Id = id;

                var document = new
                {
                    ids = new[] { id },
                    documents = new[] { $"Projeto: {newProject.Nome}" },
                    metadatas = new[] { new Dictionary<string, object>
                    {
                        ["id"] = id,
                        ["nome"] = newProject.Nome,
                        ["descricao"] = newProject.Descricao,
                        ["created_at"] = DateTime.UtcNow.ToString("O"),
                        ["entity_type"] = "project"
                    }},
                    embeddings = new[] { new double[384] } // Vector dummy
                };

                var json = JsonSerializer.Serialize(document);
                var content = new StringContent(json, Encoding.UTF8, "application/json");

                var response = await _httpClient.PostAsync($"/api/v2/collections/{_collectionName}/add", content);

                if (!response.IsSuccessStatusCode)
                {
                    var error = await response.Content.ReadAsStringAsync();
                    throw new Exception($"Erro ao criar projeto no Chroma: {error}");
                }
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Erro ao criar projeto no Chroma");
                throw;
            }
        }

        public async Task UpdateAsync(string id, Project updatedProject)
        {
            try
            {
                // Chroma não tem update direto, então fazemos upsert
                updatedProject.Id = id;

                var document = new
                {
                    ids = new[] { id },
                    documents = new[] { $"Projeto: {updatedProject.Nome}" },
                    metadatas = new[] { new Dictionary<string, object>
                    {
                        ["id"] = id,
                        ["nome"] = updatedProject.Nome,
                        ["descricao"] = updatedProject.Descricao,
                        ["updated_at"] = DateTime.UtcNow.ToString("O"),
                        ["entity_type"] = "project"
                    }},
                    embeddings = new[] { new double[384] }
                };

                var json = JsonSerializer.Serialize(document);
                var content = new StringContent(json, Encoding.UTF8, "application/json");

                var response = await _httpClient.PostAsync($"/api/v2/collections/{_collectionName}/upsert", content);

                if (!response.IsSuccessStatusCode)
                {
                    var error = await response.Content.ReadAsStringAsync();
                    throw new Exception($"Erro ao atualizar projeto no Chroma: {error}");
                }
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Erro ao atualizar projeto {Id} no Chroma", id);
                throw;
            }
        }

        public async Task SaveAsync(Project project)
        {
            try
            {
                if (string.IsNullOrEmpty(project.Id))
                {
                    await CreateAsync(project);
                }
                else
                {
                    var existing = await GetAsync(project.Id);
                    if (existing == null)
                    {
                        await CreateAsync(project);
                    }
                    else
                    {
                        await UpdateAsync(project.Id, project);
                    }
                }
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Erro ao salvar projeto no Chroma");
                throw;
            }
        }

        public async Task RemoveAsync(string id)
        {
            try
            {
                var deleteRequest = new
                {
                    ids = new[] { id }
                };

                var json = JsonSerializer.Serialize(deleteRequest);
                var content = new StringContent(json, Encoding.UTF8, "application/json");

                var response = await _httpClient.PostAsync($"/api/v2/collections/{_collectionName}/delete", content);

                if (!response.IsSuccessStatusCode)
                {
                    var error = await response.Content.ReadAsStringAsync();
                    _logger.LogWarning("Erro ao remover projeto {Id} do Chroma: {Error}", id, error);
                }
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Erro ao remover projeto {Id} do Chroma", id);
                throw;
            }
        }

        private async Task<string[]> GetCollectionsAsync()
        {
            try
            {
                var response = await _httpClient.GetAsync("/api/v2/collections");
                if (!response.IsSuccessStatusCode)
                {
                    return Array.Empty<string>();
                }

                var content = await response.Content.ReadAsStringAsync();
                var collections = JsonSerializer.Deserialize<string[]>(content);
                return collections ?? Array.Empty<string>();
            }
            catch
            {
                return Array.Empty<string>();
            }
        }

        private async Task CreateProjectsCollection()
        {
            var collection = new
            {
                name = _collectionName,
                metadata = new
                {
                    description = "Projects Collection",
                    created_at = DateTime.UtcNow.ToString("O")
                }
            };

            var json = JsonSerializer.Serialize(collection);
            var content = new StringContent(json, Encoding.UTF8, "application/json");

            var response = await _httpClient.PostAsync("/api/v2/collections", content);

            if (!response.IsSuccessStatusCode)
            {
                var error = await response.Content.ReadAsStringAsync();
                throw new Exception($"Erro ao criar collection de projetos: {error}");
            }

            _logger.LogInformation("Collection de projetos '{CollectionName}' criada no Chroma", _collectionName);
        }
    }

    public class ChromaGetResult
    {
        public string[] ids { get; set; } = Array.Empty<string>();
        public string[] documents { get; set; } = Array.Empty<string>();
        public Dictionary<string, object>[]? metadatas { get; set; }
    }
}