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

namespace RelayClient.Crypto;

/// <summary>
/// Client-side mirror of RelayServer.Services.Crypto.E2EeHelper. Identical algorithms +
/// key formats so blobs round-trip cleanly between server and client.
/// See the server class for full algorithm details.
/// </summary>
public static class E2EeHelper
{
    /// <summary>Generates a fresh RSA-2048 keypair. Called once per user on first launch and persisted via KeyStorage.</summary>
    public static (string publicKey, string privateKey) GenerateRsaKeyPair()
    {
        using var rsa = RSA.Create(2048);

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

    /// <summary>
    /// Hybrid encrypts a plaintext string for a specific recipient: fresh AES-256 key encrypts
    /// the payload (AES-GCM), then RSA-OAEP-SHA256 wraps the AES key with the recipient's
    /// public key. Returns base64-encoded fields ready to ship in a SocketEncryptedMessage.
    /// </summary>
    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)
        };
    }

    /// <summary>
    /// Reverse of EncryptForRecipient: RSA-decrypt the AES key with the recipient's private
    /// key, then AES-GCM-decrypt the ciphertext. Throws on tampered/corrupt payloads
    /// (auth tag mismatch). Returns the original UTF-8 plaintext string.
    /// </summary>
    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);
    }
}

/// <summary>The 4-tuple ciphertext bundle. Same shape on both client and server; matches SocketEncryptedMessage's encrypted fields.</summary>
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; }
}