using System.Buffers.Binary;
using System.Globalization;
using NSec.Cryptography;

namespace DodoSSH.Crypto.Tests;

/// <summary>
/// The user secret bundle and the passphrase-derived keys that wrap it. See docs/crypto.md §3.
/// </summary>
/// <remarks>
/// The encoding is checked against an independent implementation written here rather than by
/// round-tripping production code against itself. A round trip passes just as happily when both
/// directions are wrong in the same way, and this format cannot be changed after a single bundle has
/// been stored — the server holds no keys, so only clients could ever re-encrypt.
/// </remarks>
public sealed class UserSecretBundleTests
{
    private static readonly Guid Alice = Guid.Parse("0192f0c8-1a2b-7c3d-8e4f-5a6b7c8d9e0f");
    private static readonly Guid Bob = Guid.Parse("0192f0c8-1a2b-7c3d-8e4f-5a6b7c8d9e10");
    private static readonly DateTimeOffset CreatedAt = DateTimeOffset.FromUnixTimeMilliseconds(1_750_000_000_123);

    /// <summary>
    /// The layout of §3.1, pinned. Built from the fixed keys below at generation 1 and
    /// <see cref="CreatedAt"/>.
    /// </summary>
    /// <remarks>
    /// This constant is the golden vector for the bundle format. A change here that is not a
    /// deliberate, versioned format change makes every stored wrap unopenable, with no server-side
    /// remedy and no rollback.
    /// </remarks>
    private const string ExpectedEncodingHex =
        "647368312f62756e646c652f7631"          // "dsh1/bundle/v1"
        + "0001"                                 // version 1
        + "00000001"                             // key generation 1
        + "000001977420dc7b"                     // created at, Unix ms
        + "4041424344454647484950515253545556575859606162636465666768697071"  // x25519 sk
        + "8081828384858687888990919293949596979899a0a1a2a3a4a5a6a7a8a9b0b1"; // ed25519 sk

    private static byte[] EncryptionPrivateKey { get; } =
        Convert.FromHexString("4041424344454647484950515253545556575859606162636465666768697071");

    private static byte[] SigningPrivateKey { get; } =
        Convert.FromHexString("8081828384858687888990919293949596979899a0a1a2a3a4a5a6a7a8a9b0b1");

    // ---- Encoding ----

    [Fact]
    public void TheEncodedLength_IsNinetyTwoBytes()
    {
        UserSecretBundle.EncodedLength.ShouldBe(92);
        ExpectedEncodingHex.Length.ShouldBe(UserSecretBundle.EncodedLength * 2);
    }

    [Fact]
    public void TheDecoder_AcceptsAnIndependentlyBuiltEncoding()
    {
        // Proves the reader matches the specification rather than matching our writer.
        Span<byte> kek = stackalloc byte[CryptoSpec.SymmetricKeySize];
        FillKek(kek);

        var descriptor = DshAad.UserSecretBundle(Alice);
        var envelope = DshCrypto.Seal(kek, Convert.FromHexString(ExpectedEncodingHex), descriptor);

        using var bundle = UserSecretBundle.TryOpenUnder(kek, envelope, descriptor);

        bundle.ShouldNotBeNull();
        bundle.KeyGeneration.ShouldBe(1u);
        bundle.CreatedAt.ToUnixTimeMilliseconds().ShouldBe(CreatedAt.ToUnixTimeMilliseconds());
        bundle.EncryptionPublicKey.ShouldBe(PublicKeyOf(KeyAgreementAlgorithm.X25519, EncryptionPrivateKey));
        bundle.SigningPublicKey.ShouldBe(PublicKeyOf(SignatureAlgorithm.Ed25519, SigningPrivateKey));
    }

    [Fact]
    public void TheEncoder_ProducesWhatAnIndependentDecoderExpects()
    {
        Span<byte> kek = stackalloc byte[CryptoSpec.SymmetricKeySize];
        FillKek(kek);

        using var bundle = UserSecretBundle.Create(CreatedAt, keyGeneration: 3);
        var descriptor = DshAad.UserSecretBundle(Alice, keyGeneration: 3);

        var envelope = bundle.WrapUnder(kek, descriptor);
        var plaintext = DshCrypto.Open(kek, envelope, descriptor);

        plaintext.ShouldNotBeNull();
        plaintext.Length.ShouldBe(UserSecretBundle.EncodedLength);

        var decoded = DecodeIndependently(plaintext);

        decoded.Label.ShouldBe("dsh1/bundle/v1");
        decoded.Version.ShouldBe(1);
        decoded.KeyGeneration.ShouldBe(3u);
        decoded.CreatedAtUnixMilliseconds.ShouldBe(CreatedAt.ToUnixTimeMilliseconds());

        // The private keys in the encoding must be the ones whose public halves the bundle reports.
        PublicKeyOf(KeyAgreementAlgorithm.X25519, decoded.EncryptionPrivateKey)
            .ShouldBe(bundle.EncryptionPublicKey);
        PublicKeyOf(SignatureAlgorithm.Ed25519, decoded.SigningPrivateKey)
            .ShouldBe(bundle.SigningPublicKey);
    }

    [Fact]
    public void ASubMillisecondTimestamp_IsTruncatedSoTheEncodingIsStable()
    {
        using var exact = UserSecretBundle.Create(CreatedAt);
        using var noisy = UserSecretBundle.Create(CreatedAt.AddTicks(9_999));

        noisy.CreatedAt.ShouldBe(exact.CreatedAt);
    }

    // ---- The keys inside actually work ----

    [Fact]
    public void TheSigningKey_ProducesSignaturesItsPublicHalfVerifies()
    {
        using var bundle = UserSecretBundle.Create(CreatedAt);

        var statement = KeyStatementCodec.Encode(new KeyStatementFields(
            1, "https://idp.example", "alice", null,
            bundle.EncryptionPublicKey, bundle.SigningPublicKey,
            1, CreatedAt, "laptop"));

        var signature = DshSignatures.SignKeyStatement(bundle.SigningKey, statement);

        DshSignatures.VerifyKeyStatement(bundle.SigningPublicKey, statement, signature).ShouldBeTrue();
    }

    [Fact]
    public void TheEncryptionKey_OpensWhatWasSealedToItsPublicHalf()
    {
        using var bundle = UserSecretBundle.Create(CreatedAt);

        var vaultKey = VaultKeys.Create();
        var grant = VaultKeys.WrapTo(vaultKey, bundle.EncryptionPublicKey, Alice, 1);

        VaultKeys.TryUnwrap(bundle.EncryptionKey, grant, Alice, 1).ShouldBe(vaultKey);
    }

    [Fact]
    public void TheTwoKeys_AreDistinct()
    {
        using var bundle = UserSecretBundle.Create(CreatedAt);

        bundle.EncryptionPublicKey.ShouldNotBe(bundle.SigningPublicKey);
    }

    // ---- The AAD binding ----

    [Fact]
    public void AWrapForAnotherUser_DoesNotOpen()
    {
        // The property that matters most. Without it a server could hand Bob's wrap to Alice, and if
        // they ever shared a passphrase she would silently unlock his identity.
        Span<byte> kek = stackalloc byte[CryptoSpec.SymmetricKeySize];
        FillKek(kek);

        using var bundle = UserSecretBundle.Create(CreatedAt);
        var envelope = bundle.WrapUnder(kek, DshAad.UserSecretBundle(Alice));

        UserSecretBundle.TryOpenUnder(kek, envelope, DshAad.UserSecretBundle(Bob)).ShouldBeNull();
    }

    [Fact]
    public void AWrapFromAnotherKeyGeneration_DoesNotOpen()
    {
        // Stops a server serving back a superseded bundle after a key rotation.
        Span<byte> kek = stackalloc byte[CryptoSpec.SymmetricKeySize];
        FillKek(kek);

        using var bundle = UserSecretBundle.Create(CreatedAt, keyGeneration: 2);
        var envelope = bundle.WrapUnder(kek, DshAad.UserSecretBundle(Alice, keyGeneration: 2));

        UserSecretBundle.TryOpenUnder(kek, envelope, DshAad.UserSecretBundle(Alice, keyGeneration: 1))
            .ShouldBeNull();
    }

    [Fact]
    public void AWrongKey_ReturnsNullRatherThanThrowing()
    {
        // The overwhelmingly common case is a mistyped passphrase, so this is a return value.
        Span<byte> kek = stackalloc byte[CryptoSpec.SymmetricKeySize];
        FillKek(kek);

        using var bundle = UserSecretBundle.Create(CreatedAt);
        var envelope = bundle.WrapUnder(kek, DshAad.UserSecretBundle(Alice));

        UserSecretBundle.TryOpenUnder(new byte[32], envelope, DshAad.UserSecretBundle(Alice))
            .ShouldBeNull();
    }

    [Fact]
    public void ASealedWrap_DoesNotOpenAsASymmetricOne()
    {
        using var device = Key.Create(KeyAgreementAlgorithm.X25519);
        using var bundle = UserSecretBundle.Create(CreatedAt);

        var descriptor = DshAad.UserSecretBundle(Alice);
        var sealedWrap = bundle.SealTo(device.PublicKey.Export(KeyBlobFormat.RawPublicKey), descriptor);

        UserSecretBundle.TryOpenUnder(new byte[32], sealedWrap, descriptor).ShouldBeNull();
    }

    // ---- Device wraps ----

    [Fact]
    public void ADeviceWrap_OpensWithTheDeviceKey()
    {
        using var device = Key.Create(KeyAgreementAlgorithm.X25519);
        using var original = UserSecretBundle.Create(CreatedAt);

        var descriptor = DshAad.UserSecretBundle(Alice);
        var wrap = original.SealTo(device.PublicKey.Export(KeyBlobFormat.RawPublicKey), descriptor);

        using var reopened = UserSecretBundle.TryOpenSealed(device, wrap, descriptor);

        reopened.ShouldNotBeNull();
        reopened.EncryptionPublicKey.ShouldBe(original.EncryptionPublicKey);
        reopened.SigningPublicKey.ShouldBe(original.SigningPublicKey);
    }

    [Fact]
    public void ADeviceWrap_DoesNotOpenWithAnotherDevicesKey()
    {
        using var device = Key.Create(KeyAgreementAlgorithm.X25519);
        using var other = Key.Create(KeyAgreementAlgorithm.X25519);
        using var bundle = UserSecretBundle.Create(CreatedAt);

        var descriptor = DshAad.UserSecretBundle(Alice);
        var wrap = bundle.SealTo(device.PublicKey.Export(KeyBlobFormat.RawPublicKey), descriptor);

        UserSecretBundle.TryOpenSealed(other, wrap, descriptor).ShouldBeNull();
    }

    [Fact]
    public void ManyWrapsOfOneBundle_AllYieldTheSameIdentity()
    {
        // The load-bearing property of the whole hierarchy: a passphrase change re-wraps ~92 bytes
        // and touches one row, because every wrap protects the same bundle.
        using var device = Key.Create(KeyAgreementAlgorithm.X25519);
        using var bundle = UserSecretBundle.Create(CreatedAt);
        var descriptor = DshAad.UserSecretBundle(Alice);

        Span<byte> passphraseKek = stackalloc byte[CryptoSpec.SymmetricKeySize];
        FillKek(passphraseKek);

        Span<byte> recoveryKek = stackalloc byte[CryptoSpec.SymmetricKeySize];
        recoveryKek.Fill(0x5A);

        using var viaPassphrase = UserSecretBundle.TryOpenUnder(
            passphraseKek, bundle.WrapUnder(passphraseKek, descriptor), descriptor);
        using var viaRecovery = UserSecretBundle.TryOpenUnder(
            recoveryKek, bundle.WrapUnder(recoveryKek, descriptor), descriptor);
        using var viaDevice = UserSecretBundle.TryOpenSealed(
            device,
            bundle.SealTo(device.PublicKey.Export(KeyBlobFormat.RawPublicKey), descriptor),
            descriptor);

        foreach (var opened in new[] { viaPassphrase, viaRecovery, viaDevice })
        {
            opened.ShouldNotBeNull();
            opened.EncryptionPublicKey.ShouldBe(bundle.EncryptionPublicKey);
            opened.SigningPublicKey.ShouldBe(bundle.SigningPublicKey);
        }
    }

    // ---- Malformed plaintext ----

    [Theory]
    [InlineData(0)]
    [InlineData(91)]
    [InlineData(93)]
    [InlineData(200)]
    public void APlaintextOfTheWrongLength_IsRejected(int length)
    {
        Span<byte> kek = stackalloc byte[CryptoSpec.SymmetricKeySize];
        FillKek(kek);

        var descriptor = DshAad.UserSecretBundle(Alice);
        var envelope = DshCrypto.Seal(kek, new byte[length], descriptor);

        UserSecretBundle.TryOpenUnder(kek, envelope, descriptor).ShouldBeNull();
    }

    [Fact]
    public void APlaintextWithTheWrongLabel_IsRejected()
    {
        Span<byte> kek = stackalloc byte[CryptoSpec.SymmetricKeySize];
        FillKek(kek);

        var tampered = Convert.FromHexString(ExpectedEncodingHex);
        tampered[0] ^= 0xFF;

        var descriptor = DshAad.UserSecretBundle(Alice);
        var envelope = DshCrypto.Seal(kek, tampered, descriptor);

        UserSecretBundle.TryOpenUnder(kek, envelope, descriptor).ShouldBeNull();
    }

    [Fact]
    public void APlaintextWithAnUnknownVersion_IsRejected()
    {
        Span<byte> kek = stackalloc byte[CryptoSpec.SymmetricKeySize];
        FillKek(kek);

        var tampered = Convert.FromHexString(ExpectedEncodingHex);
        BinaryPrimitives.WriteUInt16BigEndian(tampered.AsSpan(14), 2);

        var descriptor = DshAad.UserSecretBundle(Alice);
        var envelope = DshCrypto.Seal(kek, tampered, descriptor);

        UserSecretBundle.TryOpenUnder(kek, envelope, descriptor).ShouldBeNull();
    }

    [Fact]
    public void APlaintextWithGenerationZero_IsRejected()
    {
        Span<byte> kek = stackalloc byte[CryptoSpec.SymmetricKeySize];
        FillKek(kek);

        var tampered = Convert.FromHexString(ExpectedEncodingHex);
        BinaryPrimitives.WriteUInt32BigEndian(tampered.AsSpan(16), 0);

        var descriptor = DshAad.UserSecretBundle(Alice);
        var envelope = DshCrypto.Seal(kek, tampered, descriptor);

        UserSecretBundle.TryOpenUnder(kek, envelope, descriptor).ShouldBeNull();
    }

    // ---- Lifetime ----

    [Fact]
    public void UsingADisposedBundle_Throws()
    {
        var bundle = UserSecretBundle.Create(CreatedAt);
        bundle.Dispose();

        Should.Throw<ObjectDisposedException>(() => bundle.SigningKey);
        Should.Throw<ObjectDisposedException>(() => bundle.EncryptionKey);
    }

    [Fact]
    public void DisposingTwice_IsHarmless()
    {
        var bundle = UserSecretBundle.Create(CreatedAt);

        bundle.Dispose();
        Should.NotThrow(bundle.Dispose);
    }

    [Fact]
    public void Create_RejectsGenerationZero()
    {
        Should.Throw<ArgumentOutOfRangeException>(() => UserSecretBundle.Create(CreatedAt, 0));
    }

    // ---- Independent codec, for the differential assertions above ----

    private static byte[] PublicKeyOf(Algorithm algorithm, byte[] privateKey)
    {
        using var key = Key.Import(algorithm, privateKey, KeyBlobFormat.RawPrivateKey);
        return key.PublicKey.Export(KeyBlobFormat.RawPublicKey);
    }

    private static void FillKek(Span<byte> destination)
    {
        for (var i = 0; i < destination.Length; i++)
        {
            destination[i] = (byte)(0x10 + i);
        }
    }

    /// <summary>Reads §3.1 from the document, independently of the production decoder.</summary>
    private static DecodedBundle DecodeIndependently(byte[] encoded) =>
        new(
            Label: System.Text.Encoding.ASCII.GetString(encoded, 0, 14),
            Version: BinaryPrimitives.ReadUInt16BigEndian(encoded.AsSpan(14)),
            KeyGeneration: BinaryPrimitives.ReadUInt32BigEndian(encoded.AsSpan(16)),
            CreatedAtUnixMilliseconds: BinaryPrimitives.ReadInt64BigEndian(encoded.AsSpan(20)),
            EncryptionPrivateKey: encoded[28..60],
            SigningPrivateKey: encoded[60..92]);

    private sealed record DecodedBundle(
        string Label,
        int Version,
        uint KeyGeneration,
        long CreatedAtUnixMilliseconds,
        byte[] EncryptionPrivateKey,
        byte[] SigningPrivateKey);
}

/// <summary>
/// Passphrase stretching and subkey derivation. See docs/crypto.md §2 and §3.
/// </summary>
/// <remarks>
/// Uses the cheapest profile throughout. These tests are about the shape of the hierarchy, not about
/// how long Argon2id takes; the default profile would add several seconds per case for no extra
/// coverage. Cost calibration is a measurement, recorded in docs/platform-flags.md.
/// </remarks>
public sealed class MasterKeyTests
{
    private static readonly Guid Alice = Guid.Parse("0192f0c8-1a2b-7c3d-8e4f-5a6b7c8d9e0f");
    private static readonly DateTimeOffset CreatedAt = DateTimeOffset.FromUnixTimeMilliseconds(1_750_000_000_123);
    private const string Passphrase = "correct horse battery staple";

    private static byte[] Salt { get; } =
        [.. Enumerable.Range(0, CryptoSpec.SaltSize).Select(i => (byte)(0x20 + i))];

    [Fact]
    public void ThePassphrase_UnlocksTheBundleItWrapped()
    {
        var descriptor = DshAad.UserSecretBundle(Alice);

        byte[] wrap;
        byte[] expectedSigningKey;

        using (var master = Derive())
        using (var bundle = UserSecretBundle.Create(CreatedAt))
        {
            wrap = master.WrapBundle(bundle, descriptor);
            expectedSigningKey = bundle.SigningPublicKey;
        }

        // A fresh derivation, as a later unlock on another device would do.
        using var reopenedMaster = Derive();
        using var reopened = reopenedMaster.TryOpenBundle(wrap, descriptor);

        reopened.ShouldNotBeNull();
        reopened.SigningPublicKey.ShouldBe(expectedSigningKey);
    }

    [Fact]
    public void AWrongPassphrase_ReturnsNull()
    {
        var descriptor = DshAad.UserSecretBundle(Alice);

        using var master = Derive();
        using var bundle = UserSecretBundle.Create(CreatedAt);
        var wrap = master.WrapBundle(bundle, descriptor);

        using var wrong = MasterKey.Derive("not the passphrase", Salt, Argon2Profile.RandomSecret);

        wrong.TryOpenBundle(wrap, descriptor).ShouldBeNull();
    }

    [Fact]
    public void ADifferentSalt_ProducesADifferentKey()
    {
        // Which is why the salt has to be cached locally: unlock must work offline, and fetching the
        // salt at unlock time would make an offline launch impossible.
        var descriptor = DshAad.UserSecretBundle(Alice);

        using var master = Derive();
        using var bundle = UserSecretBundle.Create(CreatedAt);
        var wrap = master.WrapBundle(bundle, descriptor);

        var otherSalt = Salt.ToArray();
        otherSalt[0] ^= 0xFF;

        using var other = MasterKey.Derive(Passphrase, otherSalt, Argon2Profile.RandomSecret);

        other.TryOpenBundle(wrap, descriptor).ShouldBeNull();
    }

    [Fact]
    public void TheLocalCacheKey_CannotOpenABundleWrap()
    {
        // Domain separation, tested through behaviour rather than by comparing derived bytes. The
        // cache and the vault live in different threat models and must not share a key.
        var descriptor = DshAad.UserSecretBundle(Alice);

        using var master = Derive();
        using var bundle = UserSecretBundle.Create(CreatedAt);
        var wrap = master.WrapBundle(bundle, descriptor);

        Span<byte> cacheKey = stackalloc byte[CryptoSpec.SymmetricKeySize];
        master.DeriveLocalCacheKey(cacheKey);

        UserSecretBundle.TryOpenUnder(cacheKey, wrap, descriptor).ShouldBeNull();
    }

    [Fact]
    public void TheLocalCacheKey_IsStableForTheSamePassphraseAndSalt()
    {
        Span<byte> first = stackalloc byte[CryptoSpec.SymmetricKeySize];
        Span<byte> second = stackalloc byte[CryptoSpec.SymmetricKeySize];

        using (var master = Derive())
        {
            master.DeriveLocalCacheKey(first);
        }

        using (var master = Derive())
        {
            master.DeriveLocalCacheKey(second);
        }

        first.SequenceEqual(second).ShouldBeTrue();
    }

    [Fact]
    public void ASaltShorterThanTheSpecifiedMinimum_IsRejected()
    {
        Should.Throw<ArgumentException>(() =>
            MasterKey.Derive(Passphrase, new byte[CryptoSpec.SaltSize - 1], Argon2Profile.RandomSecret));
    }

    [Theory]
    [InlineData("")]
    [InlineData(null)]
    public void AnEmptyPassphrase_IsRejected(string? passphrase)
    {
        Should.Throw<ArgumentException>(() =>
            MasterKey.Derive(passphrase!, Salt, Argon2Profile.RandomSecret));
    }

    [Fact]
    public void UsingADisposedMasterKey_Throws()
    {
        var master = Derive();
        master.Dispose();

        Should.Throw<ObjectDisposedException>(() =>
        {
            var buffer = new byte[CryptoSpec.SymmetricKeySize];
            master.DeriveLocalCacheKey(buffer);
        });
    }

    [Fact]
    public void ASubkeyBufferOfTheWrongSize_IsRejected()
    {
        using var master = Derive();

        Should.Throw<ArgumentException>(() => master.DeriveLocalCacheKey(new byte[16]));
    }

    [Fact]
    public void TheCheapestProfile_IsStillTheSpecifiedFloor()
    {
        // Guards the constant these tests lean on: if RandomSecret ever dropped below the server's
        // enrollment floor, the tests would be exercising parameters the API rejects.
        Argon2Profile.RandomSecret.MemoryKibibytes
            .ShouldBeGreaterThanOrEqualTo(
                64 * 1024,
                string.Create(CultureInfo.InvariantCulture, $"64 MiB is the enrollment floor."));
    }

    private static MasterKey Derive() =>
        MasterKey.Derive(Passphrase, Salt, Argon2Profile.RandomSecret);
}