CodeAnalizer/CodeBase/Repositories/GraphRepository.cs

using CodeBase.Models;
using Neo4j.Driver;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;

public class GraphRepository
{
    private readonly IDriver _driver;

    public GraphRepository(IDriver driver)
    {
        _driver = driver;
    }

    // ==========================================
    // 0. ИНИЦИАЛИЗАЦИЯ (Создаем пространство для векторов)
    // ==========================================
    public async Task InitializeDbAsync()
    {
        await using var session = _driver.AsyncSession();

        await session.ExecuteWriteAsync(async tx =>
        {
            // Говорим базе: "Создай индекс для поиска по сходству, если его еще нет. 
            // Размер вектора 768 (GraphCodeBERT), алгоритм - косинусное расстояние"
            await tx.RunAsync(@"
                CREATE VECTOR INDEX code_vectors IF NOT EXISTS
                FOR (c:CodeChunk) ON (c.vector)
                OPTIONS { indexConfig: {
                    `vector.dimensions`: 768,
                    `vector.similarity_function`: 'cosine'
                }}"
            );
        });
    }

    // ==========================================
    // 1. СОХРАНЕНИЕ УЗЛОВ И ВЕКТОРОВ
    // ==========================================
    public async Task SaveChunksAsync(string projectName, List<CodeChunk> chunks)
    {
        await using var session = _driver.AsyncSession();

        // 1. Делаем данные "безопасными" для базы
        var parameters = chunks.Where(c => c.Vector != null).Select(chunk => new
        {
            id = Guid.NewGuid().ToString(),
            projectName = projectName,
            filePath = chunk.FilePath,
            methodName = chunk.MethodName,
            content = chunk.Content,
            vector = chunk.Vector,
            // Защита от null и пустых строк (база их не переварит в цикле FOREACH)
            outgoingCalls = chunk.OutgoingCalls != null
                ? chunk.OutgoingCalls.Where(x => !string.IsNullOrWhiteSpace(x)).ToList()
                : new List<string>()
        }).ToList();

        await session.ExecuteWriteAsync(async tx =>
        {
            await tx.RunAsync(@"
            UNWIND $batch AS chunk
            
            // Создаем или обновляем основной метод
            MERGE (c:CodeChunk { methodName: chunk.methodName })
            SET c.id = chunk.id, 
                c.projectName = chunk.projectName, 
                c.filePath = chunk.filePath, 
                c.content = chunk.content,
                c.vector = chunk.vector
            
            // Рисуем связи только для валидных вызовов
            FOREACH (calledMethod IN chunk.outgoingCalls |
                MERGE (target:CodeChunk { methodName: calledMethod })
                MERGE (c)-[:CALLS]->(target)
            )
        ", new { batch = parameters });
        });
    }

    // ==========================================
    // 2. ПОИСК (Пока только по вектору, связи добавим позже)
    // ==========================================
    public async Task<List<GraphNodeContext>> SearchAsync(string projectName, float[] queryVector, int topK = 15)
    {
        await using var session = _driver.AsyncSession();

        return await session.ExecuteReadAsync(async tx =>
        {
            var cursor = await tx.RunAsync(@"
            // 1. Ищем самые подходящие узлы по векторному сходству
            CALL db.index.vector.queryNodes('code_vectors', $topK, $queryVector)
            YIELD node AS c, score
            WHERE c.projectName = $projectName
            
            // 2. Ищем стрелочки ВНИЗ (Кого вызывает этот метод?)
            // OPTIONAL MATCH гарантирует, что мы не потеряем узел, если связей нет
            OPTIONAL MATCH (c)-[:CALLS]->(out:CodeChunk)
            // Собираем имена вызываемых методов в массив
            WITH c, score, collect(DISTINCT out.methodName) AS outgoingDependencies
            
            // 3. Ищем стрелочки ВВЕРХ (Кто вызывает этот метод?)
            OPTIONAL MATCH (in:CodeChunk)-[:CALLS]->(c)
            WITH c, score, outgoingDependencies, collect(DISTINCT in.methodName) AS incomingDependencies
            
            // 4. Возвращаем готовую структуру для C#
            RETURN 
                c.methodName AS methodName, 
                c.filePath AS filePath, 
                c.content AS content, 
                outgoingDependencies, 
                incomingDependencies,
                score
            ORDER BY score DESC
        ", new { topK, queryVector, projectName });

            var results = new List<GraphNodeContext>();

            while (await cursor.FetchAsync())
            {
                var record = cursor.Current;
                results.Add(new GraphNodeContext
                {
                    MethodName = record["methodName"].As<string>(),
                    FilePath = record["filePath"].As<string>(),
                    Content = record["content"].As<string>(),

                    // Драйвер Neo4j возвращает массивы как IList<object>, 
                    // поэтому аккуратно кастуем их в наши списки строк
                    OutgoingDependencies = record["outgoingDependencies"].As<IList<string>>().ToList(),
                    IncomingDependencies = record["incomingDependencies"].As<IList<string>>().ToList()
                });
            }
            return results;
        });
    }
}