MoShiro/backend/src/MoShiro.Backend.Storage.Sqlite/SqliteConfigStores.cs

using System.Text.Json;
using Fleet.Core;
using Microsoft.Data.Sqlite;
using MoShiro.Backend.Application;
using MoShiro.Protocol;

namespace MoShiro.Backend.Storage.Sqlite;

public sealed class SqliteJoinTokenStore : IJoinTokenStore
{
    private readonly SqliteDatabase _db;

    public SqliteJoinTokenStore(SqliteDatabase db)
    {
        _db = db;
    }

    public ValueTask<JoinTokenInfo> CreateAsync(string id, byte[] tokenSha256, string? note, IReadOnlyList<string> labels, bool autoApprove, string kind, int maxUses, long expiresAtMs, long nowMs, string? createdBy, CancellationToken ct) =>
        _db.WriteAsync(c =>
        {
            c.Exec(
                "INSERT INTO join_tokens (id, token_sha256, note, labels, auto_approve, kind, max_uses, expires_at, created_at, created_by) VALUES (@id, @h, @note, @labels, @auto, @kind, @max, @exp, @now, @by)",
                ("@id", id), ("@h", tokenSha256), ("@note", note), ("@labels", StorageJson.ToJson(labels)), ("@auto", autoApprove ? 1 : 0), ("@kind", kind),
                ("@max", maxUses), ("@exp", expiresAtMs), ("@now", nowMs), ("@by", createdBy));
            return new JoinTokenInfo(id, note, labels, autoApprove, maxUses, 0, expiresAtMs, null, nowMs, kind, createdBy);
        }, ct);

    public ValueTask<IReadOnlyList<JoinTokenInfo>> ListAsync(CancellationToken ct) =>
        _db.ReadAsync<IReadOnlyList<JoinTokenInfo>>(
            c => c.Query(
                "SELECT id, note, labels, auto_approve, max_uses, uses, expires_at, revoked_at, created_at, kind, created_by FROM join_tokens ORDER BY created_at DESC, id",
                r => new JoinTokenInfo(r.GetString(0), r.NullableString(1), StorageJson.Strings(r.GetString(2)), r.GetInt64(3) != 0, r.GetInt32(4), r.GetInt32(5), r.GetInt64(6), r.NullableLong(7), r.GetInt64(8), r.GetString(9), r.NullableString(10))),
            ct);

    public ValueTask<JoinTokenInfo?> FindAsync(byte[] tokenSha256, CancellationToken ct) =>
        _db.ReadAsync(
            c => c.QueryOne(
                "SELECT id, note, labels, auto_approve, max_uses, uses, expires_at, revoked_at, created_at, kind, created_by FROM join_tokens WHERE token_sha256 = @h",
                r => new JoinTokenInfo(r.GetString(0), r.NullableString(1), StorageJson.Strings(r.GetString(2)), r.GetInt64(3) != 0, r.GetInt32(4), r.GetInt32(5), r.GetInt64(6), r.NullableLong(7), r.GetInt64(8), r.GetString(9), r.NullableString(10)),
                ("@h", tokenSha256)),
            ct);

    public ValueTask<bool> RevokeAsync(string id, long nowMs, CancellationToken ct) =>
        _db.WriteAsync(c => c.Exec("UPDATE join_tokens SET revoked_at = @now WHERE id = @id AND revoked_at IS NULL", ("@id", id), ("@now", nowMs)) == 1, ct);

    public ValueTask<bool> DeleteAsync(string id, CancellationToken ct) =>
        _db.WriteAsync(c => c.Exec("DELETE FROM join_tokens WHERE id = @id", ("@id", id)) == 1, ct);
}

public sealed class SqliteSettingsStore : ISettingsStore
{
    private readonly SqliteDatabase _db;

    public SqliteSettingsStore(SqliteDatabase db)
    {
        _db = db;
    }

    public ValueTask<string?> GetAsync(string key, CancellationToken ct) =>
        _db.ReadAsync(c => c.Scalar<string>("SELECT value FROM settings WHERE key = @k", ("@k", key)), ct);

    public ValueTask SetAsync(string key, string value, CancellationToken ct) =>
        _db.ExecuteAsync(c => c.Exec("INSERT INTO settings (key, value) VALUES (@k, @v) ON CONFLICT(key) DO UPDATE SET value = excluded.value", ("@k", key), ("@v", value)), ct);

    public ValueTask<IReadOnlyDictionary<string, string>> GetAllAsync(CancellationToken ct) =>
        _db.ReadAsync<IReadOnlyDictionary<string, string>>(c => c.Query("SELECT key, value FROM settings", r => (r.GetString(0), r.GetString(1))).ToDictionary(p => p.Item1, p => p.Item2), ct);

    public ValueTask<long> IncrementAsync(string key, CancellationToken ct) =>
        _db.WriteAsync(
            c => c.Scalar<long>(
                "INSERT INTO settings (key, value) VALUES (@k, '1') ON CONFLICT(key) DO UPDATE SET value = CAST(CAST(value AS INTEGER) + 1 AS TEXT) RETURNING CAST(value AS INTEGER)",
                ("@k", key)),
            ct);
}

public sealed class SqliteTaskStore : ITaskStore
{
    private const string Columns = "id, scope, agent_id, base_task_id, kind, target, params, schedule, enabled, updated_at, row_version";

    private readonly SqliteDatabase _db;

    public SqliteTaskStore(SqliteDatabase db)
    {
        _db = db;
    }

    public ValueTask<IReadOnlyList<TaskRecord>> ListAsync(CancellationToken ct) =>
        _db.ReadAsync<IReadOnlyList<TaskRecord>>(c => c.Query($"SELECT {Columns} FROM tasks ORDER BY id", Read), ct);

    public ValueTask<TaskRecord?> GetAsync(TaskId id, CancellationToken ct) =>
        _db.ReadAsync(c => c.QueryOne($"SELECT {Columns} FROM tasks WHERE id = @id", Read, ("@id", id.Value)), ct);

    public ValueTask CreateAsync(TaskRecord task, CancellationToken ct) =>
        _db.ExecuteAsync(
            c => c.Exec(
                "INSERT INTO tasks (id, scope, agent_id, base_task_id, kind, target, params, schedule, enabled, updated_at) VALUES (@id, @scope, @agent, @base, @kind, @target, @params, @schedule, @enabled, @at)",
                ("@id", task.Id.Value), ("@scope", task.Scope), ("@agent", task.AgentId?.Value), ("@base", task.BaseTaskId?.Value), ("@kind", task.Kind),
                ("@target", task.Target), ("@params", JsonSerializer.Serialize(new Dictionary<string, JsonElement>(task.Params), StorageJson.Default.ParamMap)),
                ("@schedule", task.Schedule), ("@enabled", task.Enabled ? 1 : 0), ("@at", task.UpdatedAtMs)),
            ct);

    public ValueTask<Result<TaskRecord>> UpdateAsync(TaskRecord task, long? expectedRowVersion, CancellationToken ct) =>
        _db.WriteAsync(c =>
        {
            TaskRecord? current = c.QueryOne($"SELECT {Columns} FROM tasks WHERE id = @id", Read, ("@id", task.Id.Value));
            if (current is null)
            {
                return Result<TaskRecord>.Fail(MoShiroErrors.TaskNotFound, "task not found");
            }

            if (expectedRowVersion is { } expected && expected != current.RowVersion)
            {
                return Result<TaskRecord>.Fail(MoShiroErrors.PreconditionFailed, "the task was modified by someone else");
            }

            c.Exec(
                "UPDATE tasks SET target = @target, params = @params, schedule = @schedule, enabled = @enabled, updated_at = @at, row_version = row_version + 1 WHERE id = @id",
                ("@id", task.Id.Value), ("@target", task.Target), ("@params", JsonSerializer.Serialize(new Dictionary<string, JsonElement>(task.Params), StorageJson.Default.ParamMap)),
                ("@schedule", task.Schedule), ("@enabled", task.Enabled ? 1 : 0), ("@at", task.UpdatedAtMs));
            return Result<TaskRecord>.Ok(c.QueryOne($"SELECT {Columns} FROM tasks WHERE id = @id", Read, ("@id", task.Id.Value))!);
        }, ct);

    public ValueTask DeleteAsync(TaskId id, CancellationToken ct) =>
        _db.ExecuteAsync(c => c.Exec("DELETE FROM tasks WHERE id = @id", ("@id", id.Value)), ct);

    private static TaskRecord Read(SqliteDataReader r) => new(
        new TaskId(r.GetString(0)),
        r.GetString(1),
        r.NullableString(2) is { } agent ? new NodeId(agent) : null,
        r.NullableString(3) is { } baseId ? new TaskId(baseId) : null,
        r.GetString(4),
        r.NullableString(5),
        JsonSerializer.Deserialize(r.GetString(6), StorageJson.Default.ParamMap) ?? [],
        r.GetString(7),
        r.GetInt64(8) != 0,
        r.GetInt64(9),
        r.GetInt64(10));
}

public sealed class SqliteEventStore : IEventStore
{
    private readonly SqliteDatabase _db;

    public SqliteEventStore(SqliteDatabase db)
    {
        _db = db;
    }

    public ValueTask AddAsync(NodeId? agentId, string kind, long atMs, string? detail, CancellationToken ct) =>
        _db.ExecuteAsync(c => c.Exec("INSERT INTO events (agent_id, kind, at, detail) VALUES (@a, @k, @at, @d)", ("@a", agentId?.Value), ("@k", kind), ("@at", atMs), ("@d", detail)), ct);

    public ValueTask<IReadOnlyList<EventRecord>> ListAsync(NodeId? agentId, string? kind, long fromMs, long toMs, long? beforeId, int limit, CancellationToken ct) =>
        _db.ReadAsync<IReadOnlyList<EventRecord>>(
            c => c.Query(
                "SELECT id, agent_id, kind, at, detail FROM events WHERE at >= @from AND at < @to AND (@agent IS NULL OR agent_id = @agent) AND (@kind IS NULL OR kind = @kind) " +
                "AND (@before IS NULL OR id < @before) ORDER BY id DESC LIMIT @limit",
                r => new EventRecord(r.GetInt64(0), r.NullableString(1) is { } a ? new NodeId(a) : null, r.GetString(2), r.GetInt64(3), r.NullableString(4)),
                ("@from", fromMs), ("@to", toMs), ("@agent", agentId?.Value), ("@kind", kind), ("@before", beforeId), ("@limit", limit)),
            ct);

    public ValueTask<int> DeleteOlderThanAsync(long cutoffMs, CancellationToken ct) =>
        _db.WriteAsync(c => c.Exec("DELETE FROM events WHERE at < @cutoff", ("@cutoff", cutoffMs)), ct);
}