using System.Security.Cryptography;
using System.Text;

namespace RelayServer.Services;

public static class E2EeHelper
{
    public static (string publicKey, string privateKey) GenerateRsaKeyPair()
    {
        using var rsa = RSA.Create(2048);

        return (
            Convert.ToBase64String(rsa.ExportSubjectPublicKeyInfo()),
            Convert.ToBase64String(rsa.ExportPkcs8PrivateKey())
        );
    }

    public static EncryptedPayload EncryptForRecipient(string plainText, string recipientPublicKeyBase64)
    {
        byte[] aesKey = RandomNumberGenerator.GetBytes(32);
        byte[] nonce = RandomNumberGenerator.GetBytes(12);
        byte[] plainBytes = Encoding.UTF8.GetBytes(plainText);
        byte[] cipherBytes = new byte[plainBytes.Length];
        byte[] tag = new byte[16];

        using (var aes = new AesGcm(aesKey, 16))
        {
            aes.Encrypt(nonce, plainBytes, cipherBytes, tag);
        }

        byte[] encryptedKey;
        using (var rsa = RSA.Create())
        {
            rsa.ImportSubjectPublicKeyInfo(Convert.FromBase64String(recipientPublicKeyBase64), out _);
            encryptedKey = rsa.Encrypt(aesKey, RSAEncryptionPadding.OaepSHA256);
        }

        return new EncryptedPayload
        {
            CipherText = Convert.ToBase64String(cipherBytes),
            Nonce = Convert.ToBase64String(nonce),
            Tag = Convert.ToBase64String(tag),
            EncryptedKey = Convert.ToBase64String(encryptedKey)
        };
    }

    public static string DecryptForRecipient(EncryptedPayload payload, string recipientPrivateKeyBase64)
    {
        byte[] aesKey;
        using (var rsa = RSA.Create())
        {
            rsa.ImportPkcs8PrivateKey(Convert.FromBase64String(recipientPrivateKeyBase64), out _);
            aesKey = rsa.Decrypt(Convert.FromBase64String(payload.EncryptedKey), RSAEncryptionPadding.OaepSHA256);
        }

        byte[] plainBytes = new byte[Convert.FromBase64String(payload.CipherText).Length];

        using (var aes = new AesGcm(aesKey, 16))
        {
            aes.Decrypt(
                Convert.FromBase64String(payload.Nonce),
                Convert.FromBase64String(payload.CipherText),
                Convert.FromBase64String(payload.Tag),
                plainBytes
            );
        }

        return Encoding.UTF8.GetString(plainBytes);
    }
}

public class EncryptedPayload
{
    public required string CipherText { get; set; }
    public required string Nonce { get; set; }
    public required string Tag { get; set; }
    public required string EncryptedKey { get; set; }
}