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);
    }

    public async Task 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)
        {
            await SaveMessageAsync(msg, ct);
        }
        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);
    }

    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);
}