DodoSSH/src/DodoSSH.Client.Storage/ItemStore.cs

using DodoSSH.Contracts;
using Microsoft.EntityFrameworkCore;

namespace DodoSSH.Client.Storage;

/// <summary>
/// The mirror of what the server holds.
/// </summary>
/// <remarks>
/// <para>
/// Every row here is server-confirmed state. Nothing the user has typed but not yet pushed appears in
/// this table — that lives in <see cref="OutboxStore"/>, together with the ancestor it branched from.
/// Keeping the two apart is what makes a three-way merge possible: a single table holding "the current
/// local view" would have overwritten the ancestor and left only a two-way diff, which cannot tell an
/// edit from a revert.
/// </para>
/// <para>
/// Requires an unlocked <see cref="LocalCacheProtector"/>, which is deliberate. The protected columns
/// have to be sealed on every write and opened on every read, and a store that could be constructed
/// without a key would be a store that could write one of them in the clear.
/// </para>
/// </remarks>
public sealed class ItemStore(
    IDbContextFactory<ClientCacheContext> contexts,
    LocalCacheProtector protector)
{
    /// <summary>Reads one item, tombstones included.</summary>
    public async Task<StoredItem?> FindAsync(
        Guid vaultId,
        SyncEntityType entityType,
        Guid entityId,
        CancellationToken cancellationToken)
    {
        var context = contexts.CreateDbContext();
        await using var scope = context.ConfigureAwait(false);

        var row = await context.Set<CachedItemRow>()
            .AsNoTracking()
            .SingleOrDefaultAsync(
                r => r.VaultId == vaultId && r.EntityType == entityType && r.EntityId == entityId,
                cancellationToken)
            .ConfigureAwait(false);

        return row is null ? null : ToStored(row);
    }

    /// <summary>Reads every item of one kind in a vault.</summary>
    /// <param name="vaultId">The vault.</param>
    /// <param name="entityType">Kind of item.</param>
    /// <param name="includeDeleted">
    /// Whether to return tombstones. The interface wants them excluded; the sync engine wants them,
    /// because a tombstone is the only record that an item it once knew about has gone.
    /// </param>
    /// <param name="cancellationToken">Cancellation token.</param>
    public async Task<IReadOnlyList<StoredItem>> ListAsync(
        Guid vaultId,
        SyncEntityType entityType,
        bool includeDeleted,
        CancellationToken cancellationToken)
    {
        var context = contexts.CreateDbContext();
        await using var scope = context.ConfigureAwait(false);

        var query = context.Set<CachedItemRow>()
            .AsNoTracking()
            .Where(r => r.VaultId == vaultId && r.EntityType == entityType);

        if (!includeDeleted)
        {
            query = query.Where(r => !r.IsDeleted);
        }

        var rows = await query
            .OrderBy(r => r.ChangeSequence)
            .ToListAsync(cancellationToken)
            .ConfigureAwait(false);

        return [.. rows.Select(ToStored)];
    }

    /// <summary>
    /// Writes the server's version of an item, creating or replacing the row.
    /// </summary>
    /// <remarks>
    /// Deliberately a blind overwrite. This is a mirror, and the server's answer is the truth about
    /// what the server holds; a local edit that must survive is in the outbox, and it is the sync
    /// engine's job to have merged it before calling this.
    /// </remarks>
    public async Task SaveAsync(StoredItem item, CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(item);

        var context = contexts.CreateDbContext();
        await using var scope = context.ConfigureAwait(false);

        var row = await context.Set<CachedItemRow>()
            .SingleOrDefaultAsync(
                r => r.VaultId == item.VaultId
                    && r.EntityType == item.EntityType
                    && r.EntityId == item.EntityId,
                cancellationToken)
            .ConfigureAwait(false);

        if (row is null)
        {
            row = new CachedItemRow
            {
                VaultId = item.VaultId,
                EntityType = item.EntityType,
                EntityId = item.EntityId,
            };

            context.Add(row);
        }

        Apply(row, item);

        await context.SaveChangesAsync(cancellationToken).ConfigureAwait(false);
    }

    /// <summary>
    /// Writes the server's version of many items at once, in one transaction, in the order given.
    /// </summary>
    /// <remarks>
    /// <para>
    /// <see cref="SaveAsync"/> for a page of pulled changes. A pull used to write each change on its own —
    /// a context, a read, a commit — so a first sync of a keychain with a few thousand items was a few
    /// thousand SQLite transactions, each one a sync to disk. This is one read for the rows that exist
    /// already and one commit for the lot.
    /// </para>
    /// <para>
    /// The same blind overwrite, and the same answer when one item appears more than once: they are applied
    /// in order to the same row, so the last one given is the one kept — exactly what writing them one by one
    /// would have left.
    /// </para>
    /// </remarks>
    public async Task SaveManyAsync(IReadOnlyCollection<StoredItem> items, CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(items);

        if (items.Count == 0)
        {
            return;
        }

        var vaultIds = items.Select(item => item.VaultId).Distinct().ToList();
        var entityIds = items.Select(item => item.EntityId).Distinct().ToList();

        var context = contexts.CreateDbContext();
        await using var scope = context.ConfigureAwait(false);

        var existing = await context.Set<CachedItemRow>()
            .Where(r => vaultIds.Contains(r.VaultId) && entityIds.Contains(r.EntityId))
            .ToListAsync(cancellationToken)
            .ConfigureAwait(false);

        var rows = existing.ToDictionary(row => (row.VaultId, row.EntityType, row.EntityId));

        foreach (var item in items)
        {
            if (!rows.TryGetValue((item.VaultId, item.EntityType, item.EntityId), out var row))
            {
                row = new CachedItemRow
                {
                    VaultId = item.VaultId,
                    EntityType = item.EntityType,
                    EntityId = item.EntityId,
                };

                context.Add(row);
                rows.Add((item.VaultId, item.EntityType, item.EntityId), row);
            }

            Apply(row, item);
        }

        await context.SaveChangesAsync(cancellationToken).ConfigureAwait(false);
    }

    /// <summary>Marks the items a pull from nothing has just read as ones the server still has.</summary>
    /// <remarks>One update per type in the page; see <see cref="SweepAsync"/>.</remarks>
    public async Task MarkSeenAsync(
        Guid vaultId,
        long mark,
        IReadOnlyCollection<(SyncEntityType Type, Guid Id)> seen,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(seen);

        if (seen.Count == 0)
        {
            return;
        }

        var context = contexts.CreateDbContext();
        await using var scope = context.ConfigureAwait(false);

        foreach (var group in seen.GroupBy(item => item.Type))
        {
            var type = group.Key;
            var ids = group.Select(item => item.Id).Distinct().ToList();

            await context.Set<CachedItemRow>()
                .Where(row => row.VaultId == vaultId && row.EntityType == type && ids.Contains(row.EntityId))
                .ExecuteUpdateAsync(set => set.SetProperty(row => row.SweepMark, mark), cancellationToken)
                .ConfigureAwait(false);
        }
    }

    /// <summary>The highest change sequence held for a vault, or zero for none.</summary>
    /// <remarks>Where a pull from nothing starts its sweep's upper bound; see <see cref="SweepAsync"/>.</remarks>
    public async Task<long> HighestSequenceAsync(Guid vaultId, CancellationToken cancellationToken)
    {
        var context = contexts.CreateDbContext();
        await using var scope = context.ConfigureAwait(false);

        return await context.Set<CachedItemRow>()
            .Where(row => row.VaultId == vaultId)
            .MaxAsync(row => (long?)row.ChangeSequence, cancellationToken)
            .ConfigureAwait(false) ?? 0;
    }

    /// <summary>
    /// Removes the items a completed pull from nothing did not see, which are the ones the server no longer has.
    /// </summary>
    /// <remarks>
    /// <para>
    /// The other half of the server's tombstone collection. A machine that missed a deletion for longer than
    /// the retention is refused its cursor and pulls everything again — and everything again does not include
    /// the item, so its copy here would otherwise last for ever, and its next edit would push it back.
    /// </para>
    /// <para>
    /// ◆ Two kinds of row are not evidence and are kept. One newer than <paramref name="seenUpTo"/> was
    /// written by this machine's own push during the pull: the bound is the higher of what the cache held
    /// when the pull began and what it has read since, and the server's head is past both, so a push lands
    /// above it. One with a change still waiting to go is the user's work: it is
    /// left for the push, which finds the server without the item and says so — see
    /// <c>SyncEngine.RetryAsCreateAsync</c> — rather than being thrown away here unseen.
    /// </para>
    /// </remarks>
    /// <returns>How many items went.</returns>
    public async Task<int> SweepAsync(
        Guid vaultId,
        IReadOnlyCollection<SyncEntityType> types,
        long mark,
        long seenUpTo,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(types);

        var context = contexts.CreateDbContext();
        await using var scope = context.ConfigureAwait(false);

        // Only the types the pull asked for: a row of any other type was never going to be marked.
        return await context.Set<CachedItemRow>()
            .Where(row => row.VaultId == vaultId
                && types.Contains(row.EntityType)
                && (row.SweepMark == null || row.SweepMark != mark)
                && row.ChangeSequence <= seenUpTo
                && !context.Set<OutboxRow>().Any(pending =>
                    pending.VaultId == row.VaultId
                    && pending.EntityType == row.EntityType
                    && pending.EntityId == row.EntityId))
            .ExecuteDeleteAsync(cancellationToken)
            .ConfigureAwait(false);
    }

    /// <summary>Removes a tombstone whose change has been seen by everything that needed it.</summary>
    /// <remarks>
    /// Only ever called for a row that is already a tombstone. Collecting a live item here would make
    /// it indistinguishable from one this client has never seen, and it would silently reappear on the
    /// next full pull.
    /// </remarks>
    public async Task<int> CollectTombstonesAsync(
        Guid vaultId,
        DateTimeOffset olderThan,
        CancellationToken cancellationToken)
    {
        var context = contexts.CreateDbContext();
        await using var scope = context.ConfigureAwait(false);

        return await context.Set<CachedItemRow>()
            .Where(r => r.VaultId == vaultId && r.IsDeleted && r.UpdatedAtUtc < olderThan)
            .ExecuteDeleteAsync(cancellationToken)
            .ConfigureAwait(false);
    }

    private void Apply(CachedItemRow row, StoredItem item)
    {
        row.Version = item.Version;
        row.ChangeSequence = item.ChangeSequence;
        row.IsDeleted = item.IsDeleted;
        row.UpdatedAtUtc = item.UpdatedAt;

        row.Payload = item.Payload?.Envelope;
        row.WrappedDataKey = item.Payload?.WrappedDataKey;
        row.DataKeyId = item.Payload?.DataKeyId;
        row.KeyGeneration = item.Payload?.KeyGeneration ?? 0;
        row.AadVersion = item.Payload?.AadVersion ?? 0;

        row.ProtectedFields = item.Fields is null
            ? null
            : protector.Protect(
                AadResourceTypes.For(item.EntityType),
                item.EntityId,
                PlaintextFieldsCodec.Encode(item.Fields));
    }

    private StoredItem ToStored(CachedItemRow row) =>
        new(
            row.VaultId,
            row.EntityType,
            row.EntityId,
            row.Version,
            row.ChangeSequence,
            CacheMapping.ToPayload(
                row.Payload, row.WrappedDataKey, row.DataKeyId, row.KeyGeneration, row.AadVersion),
            OpenFields(row.EntityType, row.EntityId, row.ProtectedFields),
            row.IsDeleted,
            row.UpdatedAtUtc);

    /// <remarks>
    /// A record that will not open is treated as absent rather than fatal. The cache is not the
    /// authority — a re-pull restores it — and the alternative is one stale row aborting a sync pass
    /// and stranding every change behind it.
    /// </remarks>
    private SyncPlaintextFields? OpenFields(SyncEntityType entityType, Guid entityId, byte[]? sealedFields)
    {
        if (sealedFields is null)
        {
            return null;
        }

        var plaintext = protector.TryUnprotect(
            AadResourceTypes.For(entityType), entityId, sealedFields);

        return plaintext is null ? null : PlaintextFieldsCodec.TryDecode(plaintext);
    }
}