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

namespace RelayServer.Services.Crypto;

/// <summary>
/// Hybrid RSA-2048 + AES-GCM-256 encryption. Used for any payload that needs to be
/// readable by exactly one party (the holder of a specific RSA private key).
///
///   Encrypt:
///     1. Generate a fresh 256-bit AES key and 96-bit nonce.
///     2. Encrypt the plaintext with AES-GCM → CipherText + Tag (auth tag, 128-bit).
///     3. Encrypt the AES key with the recipient's RSA public key (OAEP-SHA256).
///     4. Return all four as base64 strings in an EncryptedPayload.
///
///   Decrypt: reverse — RSA-decrypt the AES key, then AES-GCM-decrypt the ciphertext.
///
/// Why hybrid: RSA can only encrypt small inputs (~190 bytes for 2048-bit OAEP-SHA256).
/// Wrapping a symmetric key with RSA lets us encrypt arbitrarily large payloads while
/// still using the recipient's RSA keypair as the access mechanism. This is the same
/// design as PGP, TLS handshakes, etc.
///
/// The identical implementation exists in RelayClient.Crypto.E2EeHelper — they're
/// mirrored on both ends so any payload encrypted on one side decrypts on the other.
/// </summary>
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; }
}