// Reference: golang/nats-server/server/filestore.go
// Go FileStore supports two AEAD ciphers:
//   - ChaCha20-Poly1305 (StoreCipher = ChaCha, filestore.go ~line 300)
//   - AES-256-GCM        (StoreCipher = Aes,    filestore.go ~line 310)
// Both use a random 12-byte nonce prepended to the ciphertext.
// Wire format: [12:nonce][16:tag][N:ciphertext].
//
// StoreCipher and StoreCompression enums are defined here.
// FileStoreConfig.cs references them for FileStoreConfig.Cipher / .Compression.
//
// Key requirement: 32 bytes (256-bit) for both ciphers.

using System.Security.Cryptography;

namespace NATS.Server.JetStream.Storage;

// Go: server/filestore.go:85
/// <summary>
/// Selects the symmetric cipher used for block encryption.
/// Mirrors Go's StoreCipher type (filestore.go:85).
/// </summary>
public enum StoreCipher
{
    // Go: NoCipher — encryption disabled
    NoCipher = 0,

    // Go: ChaCha — ChaCha20-Poly1305
    ChaCha = 1,

    // Go: AES — AES-256-GCM
    Aes = 2,
}

// Go: server/filestore.go:106
/// <summary>
/// Selects the compression algorithm applied to message payloads.
/// Mirrors Go's StoreCompression type (filestore.go:106).
/// </summary>
public enum StoreCompression
{
    // Go: NoCompression — no compression applied
    NoCompression = 0,

    // Go: S2Compression — S2 (Snappy variant) block compression
    S2Compression = 1,
}

/// <summary>
/// Provides AEAD encrypt/decrypt operations for FileStore payloads using
/// ChaCha20-Poly1305 or AES-256-GCM, matching the Go server's encryption
/// (filestore.go ~line 300-320).
/// </summary>
internal static class AeadEncryptor
{
    /// <summary>Nonce size in bytes (96-bit / 12 bytes, standard for both ciphers).</summary>
    public const int NonceSize = 12;

    /// <summary>Authentication tag size in bytes (128-bit / 16 bytes).</summary>
    public const int TagSize = 16;

    /// <summary>Required key size in bytes (256-bit).</summary>
    public const int KeySize = 32;

    /// <summary>
    /// Encrypts <paramref name="plaintext"/> with the given <paramref name="cipher"/>
    /// and <paramref name="key"/>.
    /// </summary>
    /// <returns>
    /// Wire format: <c>[12:nonce][16:tag][N:ciphertext]</c>
    /// </returns>
    /// <exception cref="ArgumentException">If key length is not 32 bytes.</exception>
    /// <exception cref="ArgumentOutOfRangeException">If cipher is NoCipher or unknown.</exception>
    public static byte[] Encrypt(ReadOnlySpan<byte> plaintext, byte[] key, StoreCipher cipher)
    {
        ValidateKey(key);

        // Generate a random 12-byte nonce.
        var nonce = new byte[NonceSize];
        RandomNumberGenerator.Fill(nonce);

        // Output: nonce (12) + tag (16) + ciphertext (N)
        var output = new byte[NonceSize + TagSize + plaintext.Length];
        nonce.CopyTo(output.AsSpan(0, NonceSize));

        var tagDest = output.AsSpan(NonceSize, TagSize);
        var ciphertextDest = output.AsSpan(NonceSize + TagSize, plaintext.Length);

        switch (cipher)
        {
            case StoreCipher.ChaCha:
                using (var chacha = new ChaCha20Poly1305(key))
                {
                    chacha.Encrypt(nonce, plaintext, ciphertextDest, tagDest);
                }
                break;

            case StoreCipher.Aes:
                using (var aes = new AesGcm(key, TagSize))
                {
                    aes.Encrypt(nonce, plaintext, ciphertextDest, tagDest);
                }
                break;

            default:
                throw new ArgumentOutOfRangeException(nameof(cipher), cipher,
                    "Cipher must be ChaCha or Aes for AEAD encryption.");
        }

        return output;
    }

    /// <summary>
    /// Decrypts data produced by <see cref="Encrypt"/>.
    /// </summary>
    /// <returns>Plaintext bytes.</returns>
    /// <exception cref="ArgumentException">If key length is not 32 bytes or data is too short.</exception>
    /// <exception cref="CryptographicException">If authentication tag verification fails.</exception>
    public static byte[] Decrypt(ReadOnlySpan<byte> encrypted, byte[] key, StoreCipher cipher)
    {
        ValidateKey(key);

        var minLength = NonceSize + TagSize;
        if (encrypted.Length < minLength)
            throw new ArgumentException(
                $"Encrypted data is too short: {encrypted.Length} < {minLength}.",
                nameof(encrypted));

        var nonce = encrypted[..NonceSize];
        var tag = encrypted.Slice(NonceSize, TagSize);
        var ciphertext = encrypted[(NonceSize + TagSize)..];

        var plaintext = new byte[ciphertext.Length];

        switch (cipher)
        {
            case StoreCipher.ChaCha:
                using (var chacha = new ChaCha20Poly1305(key))
                {
                    chacha.Decrypt(nonce, ciphertext, tag, plaintext);
                }
                break;

            case StoreCipher.Aes:
                using (var aes = new AesGcm(key, TagSize))
                {
                    aes.Decrypt(nonce, ciphertext, tag, plaintext);
                }
                break;

            default:
                throw new ArgumentOutOfRangeException(nameof(cipher), cipher,
                    "Cipher must be ChaCha or Aes for AEAD decryption.");
        }

        return plaintext;
    }

    private static void ValidateKey(byte[] key)
    {
        if (key is null || key.Length != KeySize)
            throw new ArgumentException(
                $"Encryption key must be exactly {KeySize} bytes (got {key?.Length ?? 0}).",
                nameof(key));
    }
}
