using System.Security.Cryptography;
using ZB.MOM.WW.Secrets.Abstractions;
using ZB.MOM.WW.Secrets.Crypto;
using ZB.MOM.WW.Secrets.MasterKey;
using ZB.MOM.WW.Secrets.Tests.Fakes;

namespace ZB.MOM.WW.Secrets.Tests.Crypto;

/// <summary>
/// Secrets-007: the cipher must read the KEK id and the key from ONE resolution of the source, and
/// zero the key copy it used. Also pins the on-disk format: a row written by the pre-snapshot cipher
/// (0.6.5) must still decrypt.
/// </summary>
public class MasterKeySnapshotTests
{
    private static readonly byte[] KeyA = Enumerable.Range(0, 32).Select(i => (byte)i).ToArray();
    private static readonly byte[] KeyB = Enumerable.Range(100, 32).Select(i => (byte)i).ToArray();

    /// <summary>
    /// A file-like provider whose source is replaced after its first read: the first resolution
    /// returns key A, every later one key B (an operator swapping the key file mid-operation).
    /// </summary>
    private sealed class SwappedOnceProvider : MasterKeyProviderBase
    {
        private int _reads;

        public SwappedOnceProvider() : base(new MasterKeyOptions()) { }

        public int Reads => _reads;

        protected override byte[] ResolveKeyBytes() =>
            (byte[])(Interlocked.Increment(ref _reads) == 1 ? KeyA : KeyB).Clone();
    }

    private sealed class FixedKeyProvider : MasterKeyProviderBase
    {
        private readonly byte[] _key;

        public FixedKeyProvider(byte[] key) : base(new MasterKeyOptions()) => _key = key;

        protected override byte[] ResolveKeyBytes() => (byte[])_key.Clone();
    }

    /// <summary>Records the key buffer of every snapshot it hands out.</summary>
    private sealed class RecordingProvider(byte[] key, string kekId) : IMasterKeyProvider
    {
        public List<byte[]> Handed { get; } = new();

        public string KekId => kekId;

        public ReadOnlyMemory<byte> GetMasterKey() => (byte[])key.Clone();

        public MasterKeySnapshot GetSnapshot()
        {
            byte[] copy = (byte[])key.Clone();
            Handed.Add(copy);
            return new MasterKeySnapshot(kekId, copy);
        }
    }

    [Fact]
    public void Encrypt_stamps_the_id_of_the_key_that_actually_wrapped_the_row()
    {
        var swapping = new SwappedOnceProvider();
        var cipher = new AesGcmEnvelopeCipher(swapping);

        StoredSecret row = cipher.Encrypt(new SecretName("sql/foo"), "hunter2", SecretContentType.Text);

        // One resolution per operation, not one for the id and another for the key.
        Assert.Equal(1, swapping.Reads);

        // The row claims key A and is wrapped by key A, so a node holding key A reads it.
        var holderOfA = new AesGcmEnvelopeCipher(new FixedKeyProvider(KeyA));
        Assert.Equal(new FixedKeyProvider(KeyA).KekId, row.KekId);
        Assert.Equal("hunter2", holderOfA.Decrypt(row));
    }

    [Fact]
    public void Decrypt_and_rewrap_resolve_each_kek_once()
    {
        var aCipher = new AesGcmEnvelopeCipher(new FixedKeyProvider(KeyA));
        StoredSecret row = aCipher.Encrypt(new SecretName("sql/foo"), "hunter2", SecretContentType.Text);

        var swapping = new SwappedOnceProvider();
        Assert.Equal("hunter2", new AesGcmEnvelopeCipher(swapping).Decrypt(row));
        Assert.Equal(1, swapping.Reads);

        var oldKek = new SwappedOnceProvider();
        var newKek = new FixedKeyProvider(KeyB);
        StoredSecret rewrapped = aCipher.Rewrap(row, oldKek, newKek);
        Assert.Equal(1, oldKek.Reads);
        Assert.Equal("hunter2", new AesGcmEnvelopeCipher(new FixedKeyProvider(KeyB)).Decrypt(rewrapped));
    }

    [Fact]
    public void The_cipher_zeroes_every_key_snapshot_it_used()
    {
        var provider = new RecordingProvider(KeyA, "k1");
        var cipher = new AesGcmEnvelopeCipher(provider);

        StoredSecret row = cipher.Encrypt(new SecretName("sql/foo"), "hunter2", SecretContentType.Text);
        Assert.Equal("hunter2", cipher.Decrypt(row));
        cipher.Rewrap(row, provider, new RecordingProvider(KeyB, "k2"));

        Assert.Equal(3, provider.Handed.Count);
        Assert.All(provider.Handed, buffer => Assert.All(buffer, b => Assert.Equal(0, b)));
    }

    [Fact]
    public void The_default_snapshot_copies_so_disposing_it_never_zeroes_the_providers_own_key()
    {
        byte[] key = (byte[])KeyA.Clone();
        var provider = new FakeMasterKeyProvider("k1", key);
        var cipher = new AesGcmEnvelopeCipher(provider);

        StoredSecret row = cipher.Encrypt(new SecretName("sql/foo"), "hunter2", SecretContentType.Text);

        // The provider (which hands out its cached array) still works after the cipher zeroed its snapshot.
        Assert.Equal("hunter2", cipher.Decrypt(row));
        Assert.Equal(KeyA, provider.GetMasterKey().ToArray());
    }

    [Fact]
    public void A_row_written_by_the_0_6_5_cipher_still_decrypts()
    {
        // Generated with the 0.6.5 AesGcmEnvelopeCipher + FileMasterKeyProvider over a raw key file
        // holding bytes 0..31 (derived KEK id). The format is unchanged by the snapshot work.
        string keyPath = Path.Combine(Path.GetTempPath(), $"zb-secrets-golden-{Guid.NewGuid():N}.key");
        File.WriteAllBytes(keyPath, KeyA);
        try
        {
            var provider = new FileMasterKeyProvider(new MasterKeyOptions
            {
                Source = MasterKeySource.File,
                FilePath = keyPath,
            });
            var row = new StoredSecret
            {
                Name = new SecretName("sql/golden"),
                ContentType = SecretContentType.Text,
                Ciphertext = Convert.FromHexString("65D1B5D4FD8F538D9B58A622E202D84CDACD1F17D809DEBA0E1575DB"),
                Nonce = Convert.FromHexString("023EF785E6D7318EB0F1A238"),
                Tag = Convert.FromHexString("AAFF7A1E634C74332D4D4B75B36DE293"),
                WrappedDek = Convert.FromHexString("8BA7B8BD1A289CF97AE91386B72061C97A7CD7FF9487140590015A0ED46E79B2"),
                WrapNonce = Convert.FromHexString("04B946E3103E0E523714DDA8"),
                WrapTag = Convert.FromHexString("F542CA1BF350577B400E524B816D64A3"),
                KekId = "sha256:630dcd2966c4",
                Revision = 1,
                CreatedUtc = DateTimeOffset.UnixEpoch,
                UpdatedUtc = DateTimeOffset.UnixEpoch,
            };

            Assert.Equal("sha256:630dcd2966c4", provider.KekId);
            Assert.Equal("correct horse battery staple", new AesGcmEnvelopeCipher(provider).Decrypt(row));
        }
        finally
        {
            File.Delete(keyPath);
        }
    }

    [Fact]
    public void Disposing_a_snapshot_zeroes_its_key()
    {
        byte[] key = RandomNumberGenerator.GetBytes(32);
        var snapshot = new MasterKeySnapshot("k1", key);
        snapshot.Dispose();
        Assert.All(key, b => Assert.Equal(0, b));
    }
}