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

namespace RelayClient.Crypto;

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

        var publicKey = Convert.ToBase64String(rsa.ExportSubjectPublicKeyInfo());
        var privateKey = Convert.ToBase64String(rsa.ExportPkcs8PrivateKey());

        return (publicKey, privateKey);
    }

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

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

        var recipientPublicKey = Convert.FromBase64String(recipientPublicKeyBase64);
        byte[] encryptedKey;

        using (var rsa = RSA.Create())
        {
            rsa.ImportSubjectPublicKeyInfo(recipientPublicKey, out _);
            encryptedKey = rsa.Encrypt(aesKey, RSAEncryptionPadding.OaepSHA256);
        }

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

    public static string DecryptForRecipient(EncryptedMessagePayload payload, string recipientPrivateKeyBase64)
    {
        var encryptedKey = Convert.FromBase64String(payload.EncryptedKey);
        var privateKey = Convert.FromBase64String(recipientPrivateKeyBase64);

        byte[] aesKey;

        using (var rsa = RSA.Create())
        {
            rsa.ImportPkcs8PrivateKey(privateKey, out _);
            aesKey = rsa.Decrypt(encryptedKey, RSAEncryptionPadding.OaepSHA256);
        }

        var nonce = Convert.FromBase64String(payload.Nonce);
        var tag = Convert.FromBase64String(payload.Tag);
        var cipherBytes = Convert.FromBase64String(payload.CipherText);
        var plainBytes = new byte[cipherBytes.Length];

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

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

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