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.DefaultTypeMap.MatchNamesWithUnderscores = true;

        conn.Execute(@"
            CREATE TABLE IF NOT EXISTS sessions (
                id            INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id    TEXT NOT NULL UNIQUE,
                name          TEXT,
                created_at    TEXT NOT NULL,
                updated_at    TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS archives (
                id            INTEGER PRIMARY KEY AUTOINCREMENT,
                archive_id    TEXT NOT NULL UNIQUE,
                session_id    TEXT NOT NULL,
                name          TEXT,
                chain_count   INTEGER NOT NULL DEFAULT 0,
                message_count INTEGER NOT NULL DEFAULT 0,
                created_at    TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS chains (
                id            INTEGER PRIMARY KEY AUTOINCREMENT,
                chain_id      TEXT NOT NULL UNIQUE,
                session_id    TEXT,
                archive_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 TABLE IF NOT EXISTS agents (
                id            INTEGER PRIMARY KEY AUTOINCREMENT,
                agent_id      TEXT NOT NULL UNIQUE,
                name          TEXT NOT NULL,
                avatar        TEXT,
                description   TEXT,
                model         TEXT NOT NULL DEFAULT 'sonnet',
                persona       TEXT,
                instructions  TEXT,
                is_deleted    INTEGER NOT NULL DEFAULT 0,
                created_at    TEXT NOT NULL,
                updated_at    TEXT NOT NULL
            );

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

        // Migrations
        TryAddColumn(conn, "sessions", "agent_id", "TEXT");
        TryAddColumn(conn, "sessions", "description", "TEXT");
        TryAddColumn(conn, "sessions", "image", "TEXT");
        TryAddColumn(conn, "sessions", "workspace", "TEXT");
        TryAddColumn(conn, "sessions", "is_deleted", "INTEGER NOT NULL DEFAULT 0");
        TryAddColumn(conn, "chains", "agent_id", "TEXT");
    }

    private void TryAddColumn(SqliteConnection conn, string table, string column, string type)
    {
        try { conn.Execute($"ALTER TABLE {table} ADD COLUMN {column} {type}"); }
        catch { /* column already exists */ }
    }

    // ── Agents ──

    public async Task<AgentRecord> CreateAgent(AgentRecord agent)
    {
        using var conn = CreateConnection();
        var now = DateTime.UtcNow.ToString("o");
        agent.AgentId = Guid.NewGuid().ToString("N");
        agent.CreatedAt = now;
        agent.UpdatedAt = now;
        await conn.ExecuteAsync(@"
            INSERT INTO agents (agent_id, name, avatar, description, model, persona, instructions, is_deleted, created_at, updated_at)
            VALUES (@AgentId, @Name, @Avatar, @Description, @Model, @Persona, @Instructions, 0, @CreatedAt, @UpdatedAt)
        ", agent);
        return agent;
    }

    public async Task<IEnumerable<AgentRecord>> GetAgents()
    {
        using var conn = CreateConnection();
        return await conn.QueryAsync<AgentRecord>(
            "SELECT * FROM agents WHERE is_deleted = 0 ORDER BY name");
    }

    public async Task<AgentRecord> GetAgent(string agentId)
    {
        using var conn = CreateConnection();
        return await conn.QuerySingleOrDefaultAsync<AgentRecord>(
            "SELECT * FROM agents WHERE agent_id = @agentId AND is_deleted = 0", new { agentId });
    }

    public async Task UpdateAgent(AgentRecord agent)
    {
        using var conn = CreateConnection();
        agent.UpdatedAt = DateTime.UtcNow.ToString("o");
        await conn.ExecuteAsync(@"
            UPDATE agents SET name = @Name, avatar = @Avatar, description = @Description,
                model = @Model, persona = @Persona, instructions = @Instructions, updated_at = @UpdatedAt
            WHERE agent_id = @AgentId
        ", agent);
    }

    public async Task DeleteAgent(string agentId)
    {
        using var conn = CreateConnection();
        await conn.ExecuteAsync(
            "UPDATE agents SET is_deleted = 1, updated_at = @now WHERE agent_id = @agentId",
            new { agentId, now = DateTime.UtcNow.ToString("o") });
    }

    // ── Sessions ──

    public async Task<SessionRecord> CreateSession(string name = null, string description = null, string image = null, string agentId = null, string workspace = null)
    {
        using var conn = CreateConnection();
        var now = DateTime.UtcNow.ToString("o");
        var session = new SessionRecord
        {
            SessionId = Guid.NewGuid().ToString("N"),
            Name = name,
            Description = description,
            Image = image,
            AgentId = agentId,
            Workspace = workspace,
            CreatedAt = now,
            UpdatedAt = now
        };
        await conn.ExecuteAsync(@"
            INSERT INTO sessions (session_id, name, description, image, agent_id, workspace, created_at, updated_at)
            VALUES (@SessionId, @Name, @Description, @Image, @AgentId, @Workspace, @CreatedAt, @UpdatedAt)
        ", session);
        return session;
    }

    public async Task<IEnumerable<SessionRecord>> GetSessions()
    {
        using var conn = CreateConnection();
        return await conn.QueryAsync<SessionRecord>(
            "SELECT * FROM sessions WHERE is_deleted = 0 ORDER BY updated_at DESC");
    }

    public async Task<SessionRecord> GetSession(string sessionId)
    {
        using var conn = CreateConnection();
        return await conn.QuerySingleOrDefaultAsync<SessionRecord>(
            "SELECT * FROM sessions WHERE session_id = @sessionId AND is_deleted = 0", new { sessionId });
    }

    public async Task UpdateSessionTimestamp(string sessionId)
    {
        using var conn = CreateConnection();
        await conn.ExecuteAsync(
            "UPDATE sessions SET updated_at = @now WHERE session_id = @sessionId",
            new { sessionId, now = DateTime.UtcNow.ToString("o") });
    }

    public async Task UpdateSession(string sessionId, System.Text.Json.JsonElement body)
    {
        using var conn = CreateConnection();
        var sets = new List<string> { "updated_at = @now" };
        var p = new Dapper.DynamicParameters();
        p.Add("sessionId", sessionId);
        p.Add("now", DateTime.UtcNow.ToString("o"));

        if (body.TryGetProperty("name", out var n)) { sets.Add("name = @name"); p.Add("name", n.GetString()); }
        if (body.TryGetProperty("description", out var d)) { sets.Add("description = @desc"); p.Add("desc", d.GetString()); }
        if (body.TryGetProperty("image", out var img)) { sets.Add("image = @image"); p.Add("image", img.GetString()); }
        if (body.TryGetProperty("agentId", out var a)) { sets.Add("agent_id = @agentId2"); p.Add("agentId2", a.GetString()); }
        if (body.TryGetProperty("workspace", out var w)) { sets.Add("workspace = @workspace"); p.Add("workspace", w.GetString()); }

        await conn.ExecuteAsync($"UPDATE sessions SET {string.Join(", ", sets)} WHERE session_id = @sessionId", p);
    }

    public async Task DeleteSession(string sessionId)
    {
        using var conn = CreateConnection();
        await conn.ExecuteAsync(
            "UPDATE sessions SET is_deleted = 1, updated_at = @now WHERE session_id = @sessionId",
            new { sessionId, now = DateTime.UtcNow.ToString("o") });
    }

    // ── Archives ──

    public async Task<ArchiveRecord> CreateArchive(string sessionId, string name = null)
    {
        using var conn = CreateConnection();

        var chainCount = await conn.ExecuteScalarAsync<int>(
            "SELECT COUNT(*) FROM chains WHERE archive_id IS NULL AND session_id = @sessionId",
            new { sessionId });
        if (chainCount == 0) return null;

        var messageCount = await conn.ExecuteScalarAsync<int>(@"
            SELECT COUNT(*) FROM messages
            WHERE chain_id IN (SELECT chain_id FROM chains WHERE archive_id IS NULL AND session_id = @sessionId)",
            new { sessionId });

        var archive = new ArchiveRecord
        {
            ArchiveId = Guid.NewGuid().ToString("N"),
            SessionId = sessionId,
            Name = name,
            ChainCount = chainCount,
            MessageCount = messageCount,
            CreatedAt = DateTime.UtcNow.ToString("o")
        };

        await conn.ExecuteAsync(@"
            INSERT INTO archives (archive_id, session_id, name, chain_count, message_count, created_at)
            VALUES (@ArchiveId, @SessionId, @Name, @ChainCount, @MessageCount, @CreatedAt)
        ", archive);

        await conn.ExecuteAsync(
            "UPDATE chains SET archive_id = @archiveId WHERE archive_id IS NULL AND session_id = @sessionId",
            new { archiveId = archive.ArchiveId, sessionId });

        return archive;
    }

    public async Task<IEnumerable<ArchiveRecord>> GetArchives(string sessionId)
    {
        using var conn = CreateConnection();
        return await conn.QueryAsync<ArchiveRecord>(
            "SELECT * FROM archives WHERE session_id = @sessionId ORDER BY id DESC",
            new { sessionId });
    }

    public async Task<IEnumerable<ChainRecord>> GetArchiveChains(string archiveId)
    {
        using var conn = CreateConnection();
        return await conn.QueryAsync<ChainRecord>(
            "SELECT * FROM chains WHERE archive_id = @archiveId ORDER BY id",
            new { archiveId });
    }

    // ── 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, archive_id, agent_id, model, system_prompt, input_tokens, output_tokens, iterations, stop_reason, duration_ms, created_at)
            VALUES (@ChainId, @SessionId, @ArchiveId, @AgentId, @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>> GetActiveChains(string sessionId)
    {
        using var conn = CreateConnection();
        return await conn.QueryAsync<ChainRecord>(
            "SELECT * FROM chains WHERE session_id = @sessionId AND archive_id IS NULL ORDER BY id DESC",
            new { sessionId });
    }

    // ── 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>> GetActiveMessagesAfterId(string sessionId, long afterId, int limit = 200)
    {
        using var conn = CreateConnection();
        return await conn.QueryAsync<MessageRecord>(@"
            SELECT m.* FROM messages m
            JOIN chains c ON m.chain_id = c.chain_id
            WHERE m.id > @afterId AND c.archive_id IS NULL AND c.session_id = @sessionId
            ORDER BY m.id LIMIT @limit",
            new { sessionId, afterId, limit });
    }
}