gott/Services/DatabaseService.cs

using System;
using System.Collections.Generic;
using System.Threading.Tasks;
using Dapper;
using Gott.Models;
using Microsoft.Data.Sqlite;

namespace Gott.Services;

public class DatabaseService
{
    private readonly string _connectionString;

    public DatabaseService(string dbPath)
    {
        _connectionString = $"Data Source={dbPath}";
        InitializeDatabase();
    }

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

    private void InitializeDatabase()
    {
        using var conn = CreateConnection();

        // Dapper: snake_case → PascalCase Mapping
        Dapper.DefaultTypeMap.MatchNamesWithUnderscores = true;
        conn.Execute(@"
            CREATE TABLE IF NOT EXISTS chains (
                id            INTEGER PRIMARY KEY AUTOINCREMENT,
                chain_id      TEXT NOT NULL UNIQUE,
                session_id    TEXT,
                model         TEXT NOT NULL,
                system_prompt TEXT,
                input_tokens  INTEGER NOT NULL DEFAULT 0,
                output_tokens INTEGER NOT NULL DEFAULT 0,
                iterations    INTEGER NOT NULL DEFAULT 0,
                stop_reason   TEXT,
                duration_ms   INTEGER NOT NULL DEFAULT 0,
                created_at    TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS messages (
                id            INTEGER PRIMARY KEY AUTOINCREMENT,
                chain_id      TEXT NOT NULL,
                role          TEXT NOT NULL,
                type          TEXT NOT NULL,
                content       TEXT,
                tool_name     TEXT,
                tool_call_id  TEXT,
                is_error      INTEGER NOT NULL DEFAULT 0,
                iteration     INTEGER NOT NULL DEFAULT 0,
                created_at    TEXT NOT NULL,
                FOREIGN KEY (chain_id) REFERENCES chains(chain_id)
            );

            CREATE INDEX IF NOT EXISTS idx_chains_session ON chains(session_id);
            CREATE INDEX IF NOT EXISTS idx_messages_chain ON messages(chain_id);
        ");
    }

    // ── Chains ──

    public async Task CreateChain(ChainRecord chain)
    {
        using var conn = CreateConnection();
        chain.CreatedAt = DateTime.UtcNow.ToString("o");
        await conn.ExecuteAsync(@"
            INSERT INTO chains (chain_id, session_id, model, system_prompt, input_tokens, output_tokens, iterations, stop_reason, duration_ms, created_at)
            VALUES (@ChainId, @SessionId, @Model, @SystemPrompt, @InputTokens, @OutputTokens, @Iterations, @StopReason, @DurationMs, @CreatedAt)
        ", chain);
    }

    public async Task UpdateChain(ChainRecord chain)
    {
        using var conn = CreateConnection();
        await conn.ExecuteAsync(@"
            UPDATE chains SET
                input_tokens = @InputTokens,
                output_tokens = @OutputTokens,
                iterations = @Iterations,
                stop_reason = @StopReason,
                duration_ms = @DurationMs
            WHERE chain_id = @ChainId
        ", chain);
    }

    public async Task<ChainRecord> GetChain(string chainId)
    {
        using var conn = CreateConnection();
        return await conn.QuerySingleOrDefaultAsync<ChainRecord>(
            "SELECT * FROM chains WHERE chain_id = @chainId", new { chainId }
        );
    }

    public async Task<IEnumerable<ChainRecord>> GetRecentChains(int limit = 50)
    {
        using var conn = CreateConnection();
        return await conn.QueryAsync<ChainRecord>(
            "SELECT * FROM chains ORDER BY id DESC LIMIT @limit", new { limit }
        );
    }

    public async Task<IEnumerable<ChainRecord>> GetSessionChains(string sessionId, int limit = 100)
    {
        using var conn = CreateConnection();
        return await conn.QueryAsync<ChainRecord>(
            "SELECT * FROM chains WHERE session_id = @sessionId ORDER BY id DESC LIMIT @limit",
            new { sessionId, limit }
        );
    }

    // ── Messages ──

    public async Task<long> SaveMessage(MessageRecord record)
    {
        using var conn = CreateConnection();
        record.CreatedAt = DateTime.UtcNow.ToString("o");
        return await conn.ExecuteScalarAsync<long>(@"
            INSERT INTO messages (chain_id, role, type, content, tool_name, tool_call_id, is_error, iteration, created_at)
            VALUES (@ChainId, @Role, @Type, @Content, @ToolName, @ToolCallId, @IsError, @Iteration, @CreatedAt);
            SELECT last_insert_rowid();
        ", record);
    }

    public async Task<IEnumerable<MessageRecord>> GetChainMessages(string chainId)
    {
        using var conn = CreateConnection();
        return await conn.QueryAsync<MessageRecord>(
            "SELECT * FROM messages WHERE chain_id = @chainId ORDER BY id", new { chainId }
        );
    }

    public async Task<IEnumerable<MessageRecord>> GetMessagesAfterId(long afterId, int limit = 200)
    {
        using var conn = CreateConnection();
        return await conn.QueryAsync<MessageRecord>(
            "SELECT * FROM messages WHERE id > @afterId ORDER BY id LIMIT @limit",
            new { afterId, limit }
        );
    }
}