TentacleBot/Services/Memory/MemoryStore.cs

using Microsoft.Data.Sqlite;
using TentacleBot.Models;

namespace TentacleBot.Services.Memory;

public class MemoryStore
{
    private readonly string _connectionString;

    public MemoryStore(DatabaseInitializer db)
    {
        _connectionString = db.ConnectionString;
    }

    public async Task<List<UserRule>> GetUserRulesAsync(string userId)
    {
        var rules = new List<UserRule>();

        using var connection = new SqliteConnection(_connectionString);
        await connection.OpenAsync();

        using var cmd = connection.CreateCommand();
        cmd.CommandText = "SELECT id, user_id, rule_text, created_at FROM user_rules WHERE user_id = @userId ORDER BY created_at";
        cmd.Parameters.AddWithValue("@userId", userId);

        using var reader = await cmd.ExecuteReaderAsync();
        while (await reader.ReadAsync())
        {
            rules.Add(new UserRule(
                reader.GetInt32(0),
                reader.GetString(1),
                reader.GetString(2),
                DateTime.Parse(reader.GetString(3))
            ));
        }

        return rules;
    }

    public async Task AddUserRuleAsync(string userId, string ruleText)
    {
        using var connection = new SqliteConnection(_connectionString);
        await connection.OpenAsync();

        using var cmd = connection.CreateCommand();
        cmd.CommandText = "INSERT INTO user_rules (user_id, rule_text) VALUES (@userId, @ruleText)";
        cmd.Parameters.AddWithValue("@userId", userId);
        cmd.Parameters.AddWithValue("@ruleText", ruleText);

        await cmd.ExecuteNonQueryAsync();
    }

    public async Task<bool> RemoveUserRuleAsync(int ruleId)
    {
        using var connection = new SqliteConnection(_connectionString);
        await connection.OpenAsync();

        using var cmd = connection.CreateCommand();
        cmd.CommandText = "DELETE FROM user_rules WHERE id = @id";
        cmd.Parameters.AddWithValue("@id", ruleId);

        return await cmd.ExecuteNonQueryAsync() > 0;
    }

    public async Task<int> SaveMemoryAsync(string userId, string content, string? category, int importance = 5)
    {
        using var connection = new SqliteConnection(_connectionString);
        await connection.OpenAsync();

        using var cmd = connection.CreateCommand();
        cmd.CommandText = @"
            INSERT INTO memories (user_id, content, category, importance)
            VALUES (@userId, @content, @category, @importance);
            SELECT last_insert_rowid();";
        cmd.Parameters.AddWithValue("@userId", userId);
        cmd.Parameters.AddWithValue("@content", content);
        cmd.Parameters.AddWithValue("@category", category ?? (object)DBNull.Value);
        cmd.Parameters.AddWithValue("@importance", importance);

        var result = await cmd.ExecuteScalarAsync();
        return Convert.ToInt32(result);
    }

    public async Task<List<MemoryEntry>> SearchMemoriesAsync(string userId, string query, int limit = 5)
    {
        var memories = new List<MemoryEntry>();

        using var connection = new SqliteConnection(_connectionString);
        await connection.OpenAsync();

        var safeQuery = query.Replace("\"", "\"\"");
        using var cmd = connection.CreateCommand();
        cmd.CommandText = @"
            SELECT m.id, m.user_id, m.content, m.category, m.importance, m.created_at
            FROM memories m
            JOIN memories_fts f ON m.id = f.rowid
            WHERE memories_fts MATCH @query AND m.user_id = @userId
            ORDER BY rank
            LIMIT @limit";
        cmd.Parameters.AddWithValue("@query", $"\"{safeQuery}\"");
        cmd.Parameters.AddWithValue("@userId", userId);
        cmd.Parameters.AddWithValue("@limit", limit);

        using var reader = await cmd.ExecuteReaderAsync();
        while (await reader.ReadAsync())
        {
            memories.Add(MapMemory(reader));

            using var updateCmd = connection.CreateCommand();
            updateCmd.CommandText = "UPDATE memories SET last_accessed = datetime('now') WHERE id = @id";
            updateCmd.Parameters.AddWithValue("@id", reader.GetInt32(0));
            await updateCmd.ExecuteNonQueryAsync();
        }

        if (memories.Count == 0)
        {
            memories = await GetRecentMemoriesAsync(userId, limit);
        }

        return memories;
    }

    public async Task<List<MemoryEntry>> GetRecentMemoriesAsync(string userId, int limit = 5)
    {
        var memories = new List<MemoryEntry>();

        using var connection = new SqliteConnection(_connectionString);
        await connection.OpenAsync();

        using var cmd = connection.CreateCommand();
        cmd.CommandText = @"
            SELECT id, user_id, content, category, importance, created_at
            FROM memories
            WHERE user_id = @userId
            ORDER BY created_at DESC
            LIMIT @limit";
        cmd.Parameters.AddWithValue("@userId", userId);
        cmd.Parameters.AddWithValue("@limit", limit);

        using var reader = await cmd.ExecuteReaderAsync();
        while (await reader.ReadAsync())
        {
            memories.Add(MapMemory(reader));
        }

        return memories;
    }

    public async Task<List<MemoryEntry>> GetImportantMemoriesAsync(string userId, int limit = 5)
    {
        var memories = new List<MemoryEntry>();

        using var connection = new SqliteConnection(_connectionString);
        await connection.OpenAsync();

        using var cmd = connection.CreateCommand();
        cmd.CommandText = @"
            SELECT id, user_id, content, category, importance, created_at
            FROM memories
            WHERE user_id = @userId
            ORDER BY importance DESC, created_at DESC
            LIMIT @limit";
        cmd.Parameters.AddWithValue("@userId", userId);
        cmd.Parameters.AddWithValue("@limit", limit);

        using var reader = await cmd.ExecuteReaderAsync();
        while (await reader.ReadAsync())
        {
            memories.Add(MapMemory(reader));
        }

        return memories;
    }

    private static MemoryEntry MapMemory(SqliteDataReader reader)
    {
        return new MemoryEntry(
            reader.GetInt32(0),
            reader.GetString(1),
            reader.GetString(2),
            reader.IsDBNull(3) ? null : reader.GetString(3),
            reader.GetInt32(4),
            DateTime.Parse(reader.GetString(5))
        );
    }
}