gott/Services/VectorSearchService.cs

using System;
using System.Text;
using System.Text.RegularExpressions;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Dapper;
using Microsoft.Data.Sqlite;

namespace Gott.Services;

/// <summary>
/// Verwaltet Message-Embeddings in SQLite und bietet semantische Suche.
/// Hybrid-Suche: Vektor-Cosine-Similarity + FTS5-Keyword-Matching mit Recency-Boost.
/// </summary>
public class VectorSearchService
{
    readonly string _connectionString;
    readonly EmbeddingService _embedding;
    readonly SemaphoreSlim _embeddingLock = new(1, 1);

    public VectorSearchService(string dbPath, EmbeddingService embedding)
    {
        _connectionString = $"Data Source={dbPath}";
        _embedding = embedding;
        InitializeSchema();
    }

    SqliteConnection CreateConnection()
    {
        var conn = new SqliteConnection(_connectionString);
        conn.Open();
        return conn;
    }

    void TryExecuteSql(SqliteConnection conn, string sql)
    {
        try { conn.Execute(sql); }
        catch { /* Column/table may already exist — idempotent migration */ }
    }

    void InitializeSchema()
    {
        using var conn = CreateConnection();
        conn.Execute(@"
            CREATE TABLE IF NOT EXISTS message_embeddings (
                id           INTEGER PRIMARY KEY AUTOINCREMENT,
                message_id   INTEGER NOT NULL UNIQUE,
                chain_id     TEXT NOT NULL,
                session_id   TEXT NOT NULL,
                agent_id     TEXT,
                role         TEXT NOT NULL,
                text_preview TEXT,
                vector       BLOB NOT NULL,
                created_at   TEXT NOT NULL
            );
            CREATE INDEX IF NOT EXISTS idx_emb_session ON message_embeddings(session_id);
            CREATE INDEX IF NOT EXISTS idx_emb_agent ON message_embeddings(agent_id);
            CREATE INDEX IF NOT EXISTS idx_emb_message ON message_embeddings(message_id);

            CREATE TABLE IF NOT EXISTS memory_file_embeddings (
                id           INTEGER PRIMARY KEY AUTOINCREMENT,
                agent_id     TEXT NOT NULL,
                file_path    TEXT NOT NULL,
                chunk_index  INTEGER NOT NULL,
                text_preview TEXT,
                vector       BLOB NOT NULL,
                updated_at   TEXT NOT NULL
            );
            CREATE INDEX IF NOT EXISTS idx_memfile_agent ON memory_file_embeddings(agent_id);
        ");

        // Schema migration: add new columns (idempotent)
        TryExecuteSql(conn, "ALTER TABLE message_embeddings ADD COLUMN full_text TEXT");
        TryExecuteSql(conn, "ALTER TABLE message_embeddings ADD COLUMN embedding_model TEXT");
        TryExecuteSql(conn, "ALTER TABLE memory_file_embeddings ADD COLUMN full_text TEXT");
        TryExecuteSql(conn, "ALTER TABLE memory_file_embeddings ADD COLUMN embedding_model TEXT");

        // FTS5 virtual tables for keyword search
        TryExecuteSql(conn, "CREATE VIRTUAL TABLE IF NOT EXISTS message_embeddings_fts USING fts5(full_text, content='', content_rowid='rowid')");
        TryExecuteSql(conn, "CREATE VIRTUAL TABLE IF NOT EXISTS memory_embeddings_fts USING fts5(full_text, content='', content_rowid='rowid')");

        Console.WriteLine("[VectorSearch] Schema initialized.");
    }

    /// <summary>
    /// Generiert und speichert ein Embedding für eine Message.
    /// Thread-safe durch SemaphoreSlim (ONNX InferenceSession ist nicht thread-safe).
    /// </summary>
    public async Task EmbedMessageAsync(long messageId, string chainId, string sessionId,
        string agentId, string role, string text)
    {
        if (!_embedding.IsAvailable || string.IsNullOrWhiteSpace(text)) return;
        if (text.Length < 10) return; // Zu kurze Texte sind nicht sinnvoll

        // Prüfe ob schon embedded
        using var conn = CreateConnection();
        var exists = await conn.ExecuteScalarAsync<int>(
            "SELECT COUNT(*) FROM message_embeddings WHERE message_id = @messageId",
            new { messageId });
        if (exists > 0) return;

        // Embedding generieren (thread-safe)
        float[] vector;
        await _embeddingLock.WaitAsync();
        try
        {
            vector = _embedding.GenerateEmbedding(text);
        }
        finally
        {
            _embeddingLock.Release();
        }

        if (vector.Length == 0) return;

        var preview = text.Length > 200 ? text.Substring(0, 200) : text;

        await conn.ExecuteAsync(@"
            INSERT OR IGNORE INTO message_embeddings
                (message_id, chain_id, session_id, agent_id, role, text_preview, vector, created_at, full_text, embedding_model)
            VALUES
                (@messageId, @chainId, @sessionId, @agentId, @role, @preview, @vector, @createdAt, @fullText, @embeddingModel)",
            new
            {
                messageId,
                chainId,
                sessionId,
                agentId = agentId ?? "",
                role,
                preview,
                vector = EmbeddingService.ToBytes(vector),
                createdAt = DateTime.UtcNow.ToString("o"),
                fullText = text,
                embeddingModel = _embedding.ModelName
            });

        // Insert into FTS5 index
        try
        {
            var lastId = await conn.ExecuteScalarAsync<long>("SELECT last_insert_rowid()");
            if (lastId > 0)
            {
                await conn.ExecuteAsync(
                    "INSERT INTO message_embeddings_fts(rowid, full_text) VALUES (@id, @fullText)",
                    new { id = lastId, fullText = text });
            }
        }
        catch { /* FTS5 insert failure is non-critical */ }
    }

    /// <summary>
    /// Semantische Suche: Findet die N ähnlichsten Messages zu einer Query.
    /// Hybrid: Vektor-Cosine-Similarity + FTS5-Keyword + Recency-Boost.
    /// </summary>
    public async Task<List<SearchResult>> SearchAsync(string query, string agentId = null,
        int topN = 10, float minScore = 0.3f)
    {
        if (!_embedding.IsAvailable) return new List<SearchResult>();

        // Query embedden
        float[] queryVec;
        await _embeddingLock.WaitAsync();
        try
        {
            queryVec = _embedding.GenerateEmbedding(query);
        }
        finally
        {
            _embeddingLock.Release();
        }

        if (queryVec.Length == 0) return new List<SearchResult>();

        // Alle Embeddings laden und Cosine Similarity berechnen
        using var conn = CreateConnection();
        string sql;
        object param;

        if (!string.IsNullOrEmpty(agentId))
        {
            sql = "SELECT id, message_id, chain_id, session_id, agent_id, role, text_preview, full_text, vector, created_at FROM message_embeddings WHERE agent_id = @agentId";
            param = new { agentId };
        }
        else
        {
            sql = "SELECT id, message_id, chain_id, session_id, agent_id, role, text_preview, full_text, vector, created_at FROM message_embeddings";
            param = new { };
        }

        var rows = await conn.QueryAsync<EmbeddingRow>(sql, param);
        var rowList = rows.ToList();

        // --- Vector scores ---
        var vectorScores = new Dictionary<long, float>(); // keyed by id
        var rowMap = new Dictionary<long, EmbeddingRow>();
        foreach (var row in rowList)
        {
            var vec = EmbeddingService.FromBytes(row.Vector);
            var score = EmbeddingService.CosineSimilarity(queryVec, vec);
            vectorScores[row.Id] = score;
            rowMap[row.Id] = row;
        }

        // --- FTS5 keyword scores ---
        var ftsScores = new Dictionary<long, float>();
        try
        {
            var sanitized = SanitizeFtsQuery(query);
            if (!string.IsNullOrWhiteSpace(sanitized))
            {
                string ftsSql;
                object ftsParam;
                if (!string.IsNullOrEmpty(agentId))
                {
                    ftsSql = @"SELECT me.id, f.rank
                               FROM message_embeddings_fts f
                               JOIN message_embeddings me ON me.rowid = f.rowid
                               WHERE message_embeddings_fts MATCH @query AND me.agent_id = @agentId
                               ORDER BY f.rank LIMIT @limit";
                    ftsParam = new { query = sanitized, agentId, limit = topN * 3 };
                }
                else
                {
                    ftsSql = @"SELECT me.id, f.rank
                               FROM message_embeddings_fts f
                               JOIN message_embeddings me ON me.rowid = f.rowid
                               WHERE message_embeddings_fts MATCH @query
                               ORDER BY f.rank LIMIT @limit";
                    ftsParam = new { query = sanitized, limit = topN * 3 };
                }

                var ftsRows = await conn.QueryAsync<FtsRow>(ftsSql, ftsParam);
                var ftsList = ftsRows.ToList();

                if (ftsList.Count > 0)
                {
                    // FTS5 rank is negative (lower = better). Normalize to 0-1.
                    var minRank = ftsList.Min(r => r.Rank);
                    var maxRank = ftsList.Max(r => r.Rank);
                    var range = maxRank - minRank;
                    foreach (var r in ftsList)
                    {
                        float normalized = range == 0 ? 1.0f : (float)((maxRank - r.Rank) / range);
                        ftsScores[r.Id] = normalized;
                        // Ensure row is in rowMap (it should be, but load if needed)
                        if (!rowMap.ContainsKey(r.Id))
                        {
                            var row = await conn.QueryFirstOrDefaultAsync<EmbeddingRow>(
                                "SELECT id, message_id, chain_id, session_id, agent_id, role, text_preview, full_text, vector, created_at FROM message_embeddings WHERE id = @id",
                                new { id = r.Id });
                            if (row != null) rowMap[r.Id] = row;
                        }
                    }
                }
            }
        }
        catch { /* FTS5 query failure — fall back to vector-only */ }

        // --- Score fusion ---
        var allIds = new HashSet<long>(vectorScores.Keys);
        foreach (var id in ftsScores.Keys) allIds.Add(id);

        var results = new List<SearchResult>();
        foreach (var id in allIds)
        {
            if (!rowMap.ContainsKey(id)) continue;
            var row = rowMap[id];

            float vecScore = vectorScores.ContainsKey(id) ? vectorScores[id] : 0f;
            float ftsScore = ftsScores.ContainsKey(id) ? ftsScores[id] : 0f;
            float fusedScore = 0.6f * vecScore + 0.4f * ftsScore;

            // Recency boost: up to 15% for messages from last 30 days
            if (DateTime.TryParse(row.CreatedAt, out var created))
            {
                double hoursSince = (DateTime.UtcNow - created).TotalHours;
                double recencyBoost = Math.Max(0, 1.0 - hoursSince / (30.0 * 24.0)) * 0.15;
                fusedScore *= (float)(1.0 + recencyBoost);
            }

            if (fusedScore >= minScore)
            {
                results.Add(new SearchResult
                {
                    MessageId = row.MessageId,
                    ChainId = row.ChainId,
                    SessionId = row.SessionId,
                    AgentId = row.AgentId,
                    Role = row.Role,
                    TextPreview = row.TextPreview,
                    FullText = row.FullText ?? row.TextPreview,
                    Score = fusedScore,
                    CreatedAt = row.CreatedAt
                });
            }
        }

        return results.OrderByDescending(r => r.Score).Take(topN).ToList();
    }

    /// <summary>
    /// Gibt die Anzahl der gespeicherten Embeddings zurück.
    /// </summary>
    public int GetEmbeddingCount()
    {
        using var conn = CreateConnection();
        return conn.ExecuteScalar<int>("SELECT COUNT(*) FROM message_embeddings");
    }

    /// <summary>
    /// Gibt die IDs aller bereits eingebetteten Messages zurück.
    /// </summary>
    public async Task<HashSet<long>> GetEmbeddedMessageIdsAsync()
    {
        using var conn = CreateConnection();
        var ids = await conn.QueryAsync<long>("SELECT message_id FROM message_embeddings");
        return new HashSet<long>(ids);
    }

    /// <summary>
    /// Generiert und speichert ein Embedding für eine Memory-Datei (oder aktualisiert es).
    /// </summary>
    public async Task EmbedMemoryFileAsync(string agentId, string filePath, string text)
    {
        if (!_embedding.IsAvailable || string.IsNullOrWhiteSpace(text)) return;
        if (text.Length < 10) return;

        // Chunks erstellen mit Overlap
        var chunks = ChunkText(text, 500, 100);

        using var conn = CreateConnection();

        // Alte Embeddings für diese Datei löschen (bei Update)
        await conn.ExecuteAsync(
            "DELETE FROM memory_file_embeddings WHERE agent_id = @agentId AND file_path = @filePath",
            new { agentId, filePath });

        // Also clean up old FTS entries for this file's chunks
        try
        {
            var oldIds = await conn.QueryAsync<long>(
                "SELECT rowid FROM memory_file_embeddings WHERE agent_id = @agentId AND file_path = @filePath",
                new { agentId, filePath });
            foreach (var oldId in oldIds)
            {
                try { await conn.ExecuteAsync("INSERT INTO memory_embeddings_fts(memory_embeddings_fts, rowid, full_text) VALUES('delete', @id, '')", new { id = oldId }); } catch { }
            }
        }
        catch { }

        for (int i = 0; i < chunks.Count; i++)
        {
            float[] vector;
            await _embeddingLock.WaitAsync();
            try { vector = _embedding.GenerateEmbedding(chunks[i]); }
            finally { _embeddingLock.Release(); }
            if (vector.Length == 0) continue;

            var preview = chunks[i].Length > 200 ? chunks[i].Substring(0, 200) : chunks[i];
            await conn.ExecuteAsync(@"
                INSERT INTO memory_file_embeddings
                    (agent_id, file_path, chunk_index, text_preview, vector, updated_at, full_text, embedding_model)
                VALUES
                    (@agentId, @filePath, @chunkIndex, @preview, @vector, @updatedAt, @fullText, @embeddingModel)",
                new
                {
                    agentId,
                    filePath,
                    chunkIndex = i,
                    preview,
                    vector = EmbeddingService.ToBytes(vector),
                    updatedAt = DateTime.UtcNow.ToString("o"),
                    fullText = chunks[i],
                    embeddingModel = _embedding.ModelName
                });

            // Insert into FTS5 index
            try
            {
                var lastId = await conn.ExecuteScalarAsync<long>("SELECT last_insert_rowid()");
                if (lastId > 0)
                {
                    await conn.ExecuteAsync(
                        "INSERT INTO memory_embeddings_fts(rowid, full_text) VALUES (@id, @fullText)",
                        new { id = lastId, fullText = chunks[i] });
                }
            }
            catch { /* FTS5 insert failure is non-critical */ }
        }
    }

    /// <summary>
    /// Sucht über Memory-Dateien eines Agents.
    /// Hybrid: Vektor-Cosine-Similarity + FTS5-Keyword + Recency-Boost.
    /// </summary>
    public async Task<List<MemorySearchResult>> SearchMemoryFilesAsync(string query, string agentId, int topN = 5, float minScore = 0.4f)
    {
        if (!_embedding.IsAvailable || string.IsNullOrEmpty(agentId)) return new List<MemorySearchResult>();

        float[] queryVec;
        await _embeddingLock.WaitAsync();
        try { queryVec = _embedding.GenerateEmbedding(query); }
        finally { _embeddingLock.Release(); }
        if (queryVec.Length == 0) return new List<MemorySearchResult>();

        using var conn = CreateConnection();
        var rows = await conn.QueryAsync<MemoryFileRow>(
            "SELECT id, agent_id, file_path, chunk_index, text_preview, full_text, vector, updated_at FROM memory_file_embeddings WHERE agent_id = @agentId",
            new { agentId });
        var rowList = rows.ToList();

        // --- Vector scores ---
        var vectorScores = new Dictionary<long, float>();
        var rowMap = new Dictionary<long, MemoryFileRow>();
        foreach (var row in rowList)
        {
            var vec = EmbeddingService.FromBytes(row.Vector);
            var score = EmbeddingService.CosineSimilarity(queryVec, vec);
            vectorScores[row.Id] = score;
            rowMap[row.Id] = row;
        }

        // --- FTS5 keyword scores ---
        var ftsScores = new Dictionary<long, float>();
        try
        {
            var sanitized = SanitizeFtsQuery(query);
            if (!string.IsNullOrWhiteSpace(sanitized))
            {
                var ftsSql = @"SELECT me.id, f.rank
                               FROM memory_embeddings_fts f
                               JOIN memory_file_embeddings me ON me.rowid = f.rowid
                               WHERE memory_embeddings_fts MATCH @query AND me.agent_id = @agentId
                               ORDER BY f.rank LIMIT @limit";
                var ftsParam = new { query = sanitized, agentId, limit = topN * 3 };

                var ftsRows = await conn.QueryAsync<FtsRow>(ftsSql, ftsParam);
                var ftsList = ftsRows.ToList();

                if (ftsList.Count > 0)
                {
                    var minRank = ftsList.Min(r => r.Rank);
                    var maxRank = ftsList.Max(r => r.Rank);
                    var range = maxRank - minRank;
                    foreach (var r in ftsList)
                    {
                        float normalized = range == 0 ? 1.0f : (float)((maxRank - r.Rank) / range);
                        ftsScores[r.Id] = normalized;
                        if (!rowMap.ContainsKey(r.Id))
                        {
                            var row = await conn.QueryFirstOrDefaultAsync<MemoryFileRow>(
                                "SELECT id, agent_id, file_path, chunk_index, text_preview, full_text, vector, updated_at FROM memory_file_embeddings WHERE id = @id",
                                new { id = r.Id });
                            if (row != null) rowMap[r.Id] = row;
                        }
                    }
                }
            }
        }
        catch { /* FTS5 query failure — fall back to vector-only */ }

        // --- Score fusion ---
        var allIds = new HashSet<long>(vectorScores.Keys);
        foreach (var id in ftsScores.Keys) allIds.Add(id);

        var results = new List<MemorySearchResult>();
        foreach (var id in allIds)
        {
            if (!rowMap.ContainsKey(id)) continue;
            var row = rowMap[id];

            float vecScore = vectorScores.ContainsKey(id) ? vectorScores[id] : 0f;
            float ftsScore = ftsScores.ContainsKey(id) ? ftsScores[id] : 0f;
            float fusedScore = 0.6f * vecScore + 0.4f * ftsScore;

            // Recency boost
            if (DateTime.TryParse(row.UpdatedAt, out var updated))
            {
                double hoursSince = (DateTime.UtcNow - updated).TotalHours;
                double recencyBoost = Math.Max(0, 1.0 - hoursSince / (30.0 * 24.0)) * 0.15;
                fusedScore *= (float)(1.0 + recencyBoost);
            }

            if (fusedScore >= minScore)
            {
                results.Add(new MemorySearchResult
                {
                    AgentId = row.AgentId,
                    FilePath = row.FilePath,
                    ChunkIndex = row.ChunkIndex,
                    TextPreview = row.TextPreview,
                    FullText = row.FullText ?? row.TextPreview,
                    Score = fusedScore,
                    UpdatedAt = row.UpdatedAt
                });
            }
        }

        return results.OrderByDescending(r => r.Score).Take(topN).ToList();
    }

    /// <summary>
    /// Löscht alle Embeddings zur Neugenerierung.
    /// </summary>
    public async Task ReEmbedAllAsync()
    {
        using var conn = CreateConnection();
        await conn.ExecuteAsync("DELETE FROM message_embeddings");
        await conn.ExecuteAsync("DELETE FROM memory_file_embeddings");
        // Reset FTS5 tables
        try { await conn.ExecuteAsync("DELETE FROM message_embeddings_fts"); } catch { }
        try { await conn.ExecuteAsync("DELETE FROM memory_embeddings_fts"); } catch { }
        Console.WriteLine("[VectorSearch] All embeddings cleared for re-embedding.");
    }

    /// <summary>
    /// Sanitizes a query string for FTS5 MATCH by removing special operators.
    /// </summary>
    static string SanitizeFtsQuery(string query)
    {
        if (string.IsNullOrWhiteSpace(query)) return "";
        // Remove FTS5 special characters: * " ( ) : ^ - + AND OR NOT NEAR
        var sanitized = Regex.Replace(query, @"[""()*:^+\-]", " ");
        sanitized = Regex.Replace(sanitized, @"\b(AND|OR|NOT|NEAR)\b", " ", RegexOptions.IgnoreCase);
        sanitized = Regex.Replace(sanitized, @"\s+", " ").Trim();
        return sanitized;
    }

    static List<string> ChunkText(string text, int chunkSize, int overlap = 0)
    {
        var chunks = new List<string>();
        if (string.IsNullOrEmpty(text)) return chunks;

        // Versuche an Absätzen zu trennen
        var paragraphs = text.Split(new[] { "\n\n", "\r\n\r\n" }, StringSplitOptions.RemoveEmptyEntries);
        var current = new StringBuilder();

        // Track paragraph positions for overlap
        var paragraphStarts = new List<int>();
        var paragraphTexts = new List<string>();
        int pos = 0;
        foreach (var para in paragraphs)
        {
            paragraphStarts.Add(pos);
            paragraphTexts.Add(para);
            pos += para.Length + 2; // approximate separator length
        }

        for (int pi = 0; pi < paragraphTexts.Count; pi++)
        {
            var para = paragraphTexts[pi];
            if (current.Length + para.Length > chunkSize && current.Length > 0)
            {
                chunks.Add(current.ToString().Trim());

                // Apply overlap: backtrack to include some text from end of current chunk
                current.Clear();
                if (overlap > 0 && pi > 0)
                {
                    // Collect paragraphs from the end until we have ~overlap chars
                    int overlapChars = 0;
                    int startPi = pi;
                    for (int bi = pi - 1; bi >= 0 && overlapChars < overlap; bi--)
                    {
                        overlapChars += paragraphTexts[bi].Length;
                        startPi = bi;
                    }
                    for (int bi = startPi; bi < pi; bi++)
                    {
                        if (current.Length > 0) current.AppendLine();
                        current.Append(paragraphTexts[bi]);
                    }
                }
            }
            if (current.Length > 0) current.AppendLine();
            current.Append(para);
        }
        if (current.Length > 0)
            chunks.Add(current.ToString().Trim());

        return chunks;
    }

    class FtsRow
    {
        public long Id { get; set; }
        public double Rank { get; set; }
    }

    public class MemoryFileRow
    {
        public long Id { get; set; }
        public string AgentId { get; set; }
        public string FilePath { get; set; }
        public int ChunkIndex { get; set; }
        public string TextPreview { get; set; }
        public string FullText { get; set; }
        public byte[] Vector { get; set; }
        public string UpdatedAt { get; set; }
    }

    public class MemorySearchResult
    {
        public string AgentId { get; set; }
        public string FilePath { get; set; }
        public int ChunkIndex { get; set; }
        public string TextPreview { get; set; }
        public string FullText { get; set; }
        public float Score { get; set; }
        public string UpdatedAt { get; set; }
    }

    public class EmbeddingRow
    {
        public long Id { get; set; }
        public long MessageId { get; set; }
        public string ChainId { get; set; }
        public string SessionId { get; set; }
        public string AgentId { get; set; }
        public string Role { get; set; }
        public string TextPreview { get; set; }
        public string FullText { get; set; }
        public byte[] Vector { get; set; }
        public string CreatedAt { get; set; }
    }

    public class SearchResult
    {
        public long MessageId { get; set; }
        public string ChainId { get; set; }
        public string SessionId { get; set; }
        public string AgentId { get; set; }
        public string Role { get; set; }
        public string TextPreview { get; set; }
        public string FullText { get; set; }
        public float Score { get; set; }
        public string CreatedAt { get; set; }
    }
}