// Reference: golang/nats-server/server/filestore.go
// Go FileStore uses ChaCha20-Poly1305 and AES-256-GCM for block encryption:
//   - StoreCipher=ChaCha → ChaCha20-Poly1305 (filestore.go ~line 300)
//   - StoreCipher=AES    → AES-256-GCM        (filestore.go ~line 310)
// Wire format: [12:nonce][16:tag][N:ciphertext]

using System.Security.Cryptography;
using NATS.Server.JetStream.Storage;

namespace NATS.Server.Tests.JetStream.Storage;

public sealed class AeadEncryptorTests
{
    // 32-byte (256-bit) test key.
    private static byte[] TestKey => "nats-aead-test-key-for-32bytes!!"u8.ToArray();

    // Go: TestFileStoreEncrypted server/filestore_test.go:4204 (ChaCha permutation)
    [Fact]
    public void ChaCha_encrypt_decrypt_round_trips()
    {
        var plaintext = "Hello, ChaCha20-Poly1305!"u8.ToArray();
        var key = TestKey;

        var encrypted = AeadEncryptor.Encrypt(plaintext, key, StoreCipher.ChaCha);
        var decrypted = AeadEncryptor.Decrypt(encrypted, key, StoreCipher.ChaCha);

        decrypted.ShouldBe(plaintext);
    }

    // Go: TestFileStoreEncrypted server/filestore_test.go:4204 (AES permutation)
    [Fact]
    public void AesGcm_encrypt_decrypt_round_trips()
    {
        var plaintext = "Hello, AES-256-GCM!"u8.ToArray();
        var key = TestKey;

        var encrypted = AeadEncryptor.Encrypt(plaintext, key, StoreCipher.Aes);
        var decrypted = AeadEncryptor.Decrypt(encrypted, key, StoreCipher.Aes);

        decrypted.ShouldBe(plaintext);
    }

    [Fact]
    public void ChaCha_empty_plaintext_round_trips()
    {
        var encrypted = AeadEncryptor.Encrypt([], TestKey, StoreCipher.ChaCha);
        var decrypted = AeadEncryptor.Decrypt(encrypted, TestKey, StoreCipher.ChaCha);
        decrypted.ShouldBeEmpty();
    }

    [Fact]
    public void AesGcm_empty_plaintext_round_trips()
    {
        var encrypted = AeadEncryptor.Encrypt([], TestKey, StoreCipher.Aes);
        var decrypted = AeadEncryptor.Decrypt(encrypted, TestKey, StoreCipher.Aes);
        decrypted.ShouldBeEmpty();
    }

    [Fact]
    public void ChaCha_encrypted_blob_has_correct_overhead()
    {
        var plaintext = new byte[100];
        var encrypted = AeadEncryptor.Encrypt(plaintext, TestKey, StoreCipher.ChaCha);

        // Expected: nonce (12) + tag (16) + ciphertext (100) = 128
        encrypted.Length.ShouldBe(AeadEncryptor.NonceSize + AeadEncryptor.TagSize + plaintext.Length);
    }

    [Fact]
    public void AesGcm_encrypted_blob_has_correct_overhead()
    {
        var plaintext = new byte[100];
        var encrypted = AeadEncryptor.Encrypt(plaintext, TestKey, StoreCipher.Aes);

        // Expected: nonce (12) + tag (16) + ciphertext (100) = 128
        encrypted.Length.ShouldBe(AeadEncryptor.NonceSize + AeadEncryptor.TagSize + plaintext.Length);
    }

    // Go: TestFileStoreRestoreEncryptedWithNoKeyFuncFails filestore_test.go:5134
    [Fact]
    public void ChaCha_wrong_key_throws_CryptographicException()
    {
        var plaintext = "secret data"u8.ToArray();
        var encrypted = AeadEncryptor.Encrypt(plaintext, TestKey, StoreCipher.ChaCha);

        var wrongKey = "wrong-key-wrong-key-wrong-key!!!"u8.ToArray();
        Should.Throw<CryptographicException>(
            () => AeadEncryptor.Decrypt(encrypted, wrongKey, StoreCipher.ChaCha));
    }

    [Fact]
    public void AesGcm_wrong_key_throws_CryptographicException()
    {
        var plaintext = "secret data"u8.ToArray();
        var encrypted = AeadEncryptor.Encrypt(plaintext, TestKey, StoreCipher.Aes);

        var wrongKey = "wrong-key-wrong-key-wrong-key!!!"u8.ToArray();
        Should.Throw<CryptographicException>(
            () => AeadEncryptor.Decrypt(encrypted, wrongKey, StoreCipher.Aes));
    }

    [Fact]
    public void ChaCha_tampered_ciphertext_throws_CryptographicException()
    {
        var plaintext = "tamper me"u8.ToArray();
        var encrypted = AeadEncryptor.Encrypt(plaintext, TestKey, StoreCipher.ChaCha);

        // Flip a bit in the ciphertext portion (after nonce+tag).
        encrypted[^1] ^= 0xFF;

        Should.Throw<CryptographicException>(
            () => AeadEncryptor.Decrypt(encrypted, TestKey, StoreCipher.ChaCha));
    }

    [Fact]
    public void AesGcm_tampered_ciphertext_throws_CryptographicException()
    {
        var plaintext = "tamper me"u8.ToArray();
        var encrypted = AeadEncryptor.Encrypt(plaintext, TestKey, StoreCipher.Aes);

        // Flip a bit in the ciphertext portion.
        encrypted[^1] ^= 0xFF;

        Should.Throw<CryptographicException>(
            () => AeadEncryptor.Decrypt(encrypted, TestKey, StoreCipher.Aes));
    }

    [Fact]
    public void ChaCha_tampered_tag_throws_CryptographicException()
    {
        var plaintext = "tamper tag"u8.ToArray();
        var encrypted = AeadEncryptor.Encrypt(plaintext, TestKey, StoreCipher.ChaCha);

        // Flip a bit in the tag (bytes 12-27).
        encrypted[AeadEncryptor.NonceSize] ^= 0xFF;

        Should.Throw<CryptographicException>(
            () => AeadEncryptor.Decrypt(encrypted, TestKey, StoreCipher.ChaCha));
    }

    [Fact]
    public void Key_shorter_than_32_bytes_throws_ArgumentException()
    {
        var shortKey = new byte[16];
        Should.Throw<ArgumentException>(
            () => AeadEncryptor.Encrypt("data"u8.ToArray(), shortKey, StoreCipher.ChaCha));
    }

    [Fact]
    public void Key_longer_than_32_bytes_throws_ArgumentException()
    {
        var longKey = new byte[64];
        Should.Throw<ArgumentException>(
            () => AeadEncryptor.Encrypt("data"u8.ToArray(), longKey, StoreCipher.ChaCha));
    }

    [Fact]
    public void Decrypt_data_too_short_throws_ArgumentException()
    {
        // Less than nonce (12) + tag (16) = 28 bytes minimum.
        var tooShort = new byte[10];
        Should.Throw<ArgumentException>(
            () => AeadEncryptor.Decrypt(tooShort, TestKey, StoreCipher.ChaCha));
    }

    [Fact]
    public void ChaCha_each_encrypt_produces_different_ciphertext()
    {
        // Nonce is random per call so ciphertexts differ even for same plaintext.
        var plaintext = "same plaintext"u8.ToArray();
        var enc1 = AeadEncryptor.Encrypt(plaintext, TestKey, StoreCipher.ChaCha);
        var enc2 = AeadEncryptor.Encrypt(plaintext, TestKey, StoreCipher.ChaCha);

        enc1.ShouldNotBe(enc2);
    }

    [Fact]
    public void ChaCha_large_payload_round_trips()
    {
        var plaintext = new byte[64 * 1024]; // 64 KB
        Random.Shared.NextBytes(plaintext);

        var encrypted = AeadEncryptor.Encrypt(plaintext, TestKey, StoreCipher.ChaCha);
        var decrypted = AeadEncryptor.Decrypt(encrypted, TestKey, StoreCipher.ChaCha);

        decrypted.ShouldBe(plaintext);
    }

    [Fact]
    public void AesGcm_large_payload_round_trips()
    {
        var plaintext = new byte[64 * 1024]; // 64 KB
        Random.Shared.NextBytes(plaintext);

        var encrypted = AeadEncryptor.Encrypt(plaintext, TestKey, StoreCipher.Aes);
        var decrypted = AeadEncryptor.Decrypt(encrypted, TestKey, StoreCipher.Aes);

        decrypted.ShouldBe(plaintext);
    }
}