Catalogger.NET/Catalogger.Backend/Database/Queries/MessageRepository.cs

using System.Text.Json;
using Catalogger.Backend.Extensions;
using Microsoft.EntityFrameworkCore;
using Remora.Discord.API.Abstractions.Gateway.Events;
using DbMessage = Catalogger.Backend.Database.Models.Message;

namespace Catalogger.Backend.Database.Queries;

public class MessageRepository(
    ILogger logger,
    DatabaseContext db,
    IEncryptionService encryptionService
)
{
    private readonly ILogger _logger = logger.ForContext<MessageRepository>();

    public async Task SaveMessageAsync(IMessageCreate msg, CancellationToken ct = default)
    {
        _logger.Debug("Saving message {MessageId}", msg.ID);

        var metadata = new Metadata(
            IsWebhook: msg.WebhookID.HasValue,
            msg.Attachments.Select(a => new Attachment(a.Filename, a.Size, a.ContentType.Value))
        );

        var dbMessage = new DbMessage
        {
            Id = msg.ID.ToUlong(),
            UserId = msg.Author.ID.ToUlong(),
            ChannelId = msg.ChannelID.ToUlong(),
            GuildId = msg.GuildID.ToUlong(),

            EncryptedContent = await Task.Run(
                () =>
                    encryptionService.Encrypt(
                        string.IsNullOrWhiteSpace(msg.Content) ? "None" : msg.Content
                    ),
                ct
            ),
            EncryptedUsername = await Task.Run(
                () => encryptionService.Encrypt(msg.Author.Tag()),
                ct
            ),
            EncryptedMetadata = await Task.Run(
                () => encryptionService.Encrypt(JsonSerializer.Serialize(metadata)),
                ct
            ),
            AttachmentSize = msg.Attachments.Select(a => a.Size).Sum(),
        };

        db.Add(dbMessage);
        await db.SaveChangesAsync(ct);
    }

    /// <summary>
    /// Updates an edited message.
    /// </summary>
    /// <returns>true if the message was already stored and got updated,
    /// false if the message wasn't stored and was newly inserted.</returns>
    public async Task<bool> UpdateMessageAsync(IMessageCreate msg, CancellationToken ct = default)
    {
        _logger.Debug("Updating message {MessageId}", msg.ID);

        var tx = await db.Database.BeginTransactionAsync(ct);
        var (isStored, _) = await HasProxyInfoAsync(msg.ID.Value);
        if (!isStored)
        {
            _logger.Debug("Edited message {MessageId} is not stored yet, storing it", msg.ID);
            await SaveMessageAsync(msg, ct);
            await tx.CommitAsync(ct);
            return false;
        }
        else
        {
            var metadata = new Metadata(
                IsWebhook: msg.WebhookID.HasValue,
                msg.Attachments.Select(a => new Attachment(a.Filename, a.Size, a.ContentType.Value))
            );

            var dbMsg = await db.Messages.FindAsync(msg.ID.Value);
            if (dbMsg == null)
                throw new CataloggerError(
                    "Message was null despite HasProxyInfoAsync returning true"
                );

            dbMsg.EncryptedContent = await Task.Run(
                () =>
                    encryptionService.Encrypt(
                        string.IsNullOrWhiteSpace(msg.Content) ? "None" : msg.Content
                    ),
                ct
            );
            dbMsg.EncryptedUsername = await Task.Run(
                () => encryptionService.Encrypt(msg.Author.Tag()),
                ct
            );
            dbMsg.EncryptedMetadata = await Task.Run(
                () => encryptionService.Encrypt(JsonSerializer.Serialize(metadata)),
                ct
            );

            db.Update(dbMsg);
            await db.SaveChangesAsync(ct);
            await tx.CommitAsync(ct);
            return true;
        }
    }

    public async Task<Message?> GetMessageAsync(ulong id, CancellationToken ct = default)
    {
        _logger.Debug("Retrieving message {MessageId}", id);

        var dbMsg = await db.Messages.FindAsync(id);
        if (dbMsg == null)
            return null;

        return new Message(
            dbMsg.Id,
            dbMsg.OriginalId,
            dbMsg.UserId,
            dbMsg.ChannelId,
            dbMsg.GuildId,
            dbMsg.Member,
            dbMsg.System,
            Username: await Task.Run(() => encryptionService.Decrypt(dbMsg.EncryptedUsername), ct),
            Content: await Task.Run(() => encryptionService.Decrypt(dbMsg.EncryptedContent), ct),
            Metadata: dbMsg.EncryptedMetadata != null
                ? JsonSerializer.Deserialize<Metadata>(
                    await Task.Run(() => encryptionService.Decrypt(dbMsg.EncryptedMetadata), ct)
                )
                : null,
            dbMsg.AttachmentSize
        );
    }

    /// <summary>
    /// Checks if a message has proxy information.
    /// If yes, returns (true, true). If no, returns (true, false). If the message isn't saved at all, returns (false, false).
    /// </summary>
    public async Task<(bool, bool)> HasProxyInfoAsync(ulong id)
    {
        _logger.Debug("Checking if message {MessageId} has proxy information", id);

        var msg = await db
            .Messages.Select(m => new { m.Id, m.OriginalId })
            .FirstOrDefaultAsync(m => m.Id == id);
        return (msg != null, msg?.OriginalId != null);
    }

    public async Task SetProxiedMessageDataAsync(
        ulong id,
        ulong originalId,
        ulong authorId,
        string? systemId,
        string? memberId
    )
    {
        _logger.Debug("Setting proxy information for message {MessageId}", id);

        var message = await db.Messages.FirstOrDefaultAsync(m => m.Id == id);
        if (message == null)
        {
            _logger.Debug("Message {MessageId} not found", id);
            return;
        }

        _logger.Debug("Updating message {MessageId}", id);

        message.OriginalId = originalId;
        message.UserId = authorId;
        message.System = systemId;
        message.Member = memberId;

        db.Update(message);
        await db.SaveChangesAsync();
    }

    public async Task<bool> IsMessageIgnoredAsync(ulong id, CancellationToken ct = default)
    {
        _logger.Debug("Checking if message {MessageId} is ignored", id);
        return await db.IgnoredMessages.FirstOrDefaultAsync(m => m.Id == id, ct) != null;
    }

    public record Message(
        ulong Id,
        ulong? OriginalId,
        ulong UserId,
        ulong ChannelId,
        ulong GuildId,
        string? Member,
        string? System,
        string Username,
        string Content,
        Metadata? Metadata,
        int AttachmentSize
    );

    public record Metadata(bool IsWebhook, IEnumerable<Attachment> Attachments);

    public record Attachment(string Filename, int Size, string ContentType);
}