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

namespace DodoSSH.Crypto;

/// <summary>
/// A user's identity key pair, in the form that gets wrapped. See docs/crypto.md §3.
/// </summary>
/// <remarks>
/// <para>
/// This is the load-bearing structural choice of the whole key hierarchy. The bundle is stored
/// server-side as N independent wraps of the <em>same</em> bytes — one per passphrase, device,
/// recovery code and escrow key — so a passphrase change re-wraps about a hundred bytes and updates
/// one row. No vault data is re-encrypted, and no other member is involved. Adding a device or a
/// recovery code is likewise one more wrap row. Encrypting vault keys under the passphrase key
/// directly would have made every one of those a fan-out over the entire vault.
/// </para>
/// <para>
/// Private keys live in libsodium's guarded, <c>mlock</c>ed allocations rather than in a
/// <see cref="byte"/> array, so they are not paged to disk and do not appear in a core dump. They
/// are created exportable, which is a deliberate trade: re-wrapping the same bundle for a passphrase
/// change or a new device requires re-encoding it, and the alternative — holding the plaintext in a
/// long-lived managed array so the keys need not be exportable — keeps the identical secret in
/// strictly worse memory. Every export here is into a buffer that is zeroed before the method
/// returns.
/// </para>
/// </remarks>
public sealed class UserSecretBundle : IDisposable
{
    /// <summary>Domain-separating prefix of the encoding.</summary>
    public static ReadOnlySpan<byte> Label => "dsh1/bundle/v1"u8;

    /// <summary>Encoding version this implementation writes.</summary>
    public const int CurrentVersion = 1;

    /// <summary>
    /// Total encoded length. Fixed, because every field is fixed-width.
    /// </summary>
    public const int EncodedLength =
        14                                  // Label
        + sizeof(ushort)                    // version
        + sizeof(uint)                      // key generation
        + sizeof(long)                      // created at, Unix milliseconds
        + (CryptoSpec.SymmetricKeySize * 2); // X25519 then Ed25519 private keys

    private const int OffsetVersion = 14;
    private const int OffsetKeyGeneration = 16;
    private const int OffsetCreatedAt = 20;
    private const int OffsetEncryptionKey = 28;
    private const int OffsetSigningKey = 60;

    private readonly Key encryptionKey;
    private readonly Key signingKey;
    private bool disposed;

    private UserSecretBundle(Key encryptionKey, Key signingKey, uint keyGeneration, DateTimeOffset createdAt)
    {
        this.encryptionKey = encryptionKey;
        this.signingKey = signingKey;
        KeyGeneration = keyGeneration;
        CreatedAt = createdAt;

        EncryptionPublicKey = encryptionKey.PublicKey.Export(KeyBlobFormat.RawPublicKey);
        SigningPublicKey = signingKey.PublicKey.Export(KeyBlobFormat.RawPublicKey);
    }

    /// <summary>Generation of this identity key pair, starting at 1.</summary>
    public uint KeyGeneration { get; }

    /// <summary>When the keys were generated, truncated to milliseconds.</summary>
    public DateTimeOffset CreatedAt { get; }

    /// <summary>X25519 public key, 32 bytes. Safe to publish.</summary>
    public byte[] EncryptionPublicKey { get; }

    /// <summary>Ed25519 public key, 32 bytes. Safe to publish.</summary>
    public byte[] SigningPublicKey { get; }

    /// <summary>The X25519 key, for unwrapping vault keys sealed to this user.</summary>
    public Key EncryptionKey => Alive().encryptionKey;

    /// <summary>The Ed25519 key, for signing key statements and grants.</summary>
    public Key SigningKey => Alive().signingKey;

    /// <summary>Generates a fresh identity key pair.</summary>
    /// <param name="createdAt">Creation timestamp; truncated to milliseconds.</param>
    /// <param name="keyGeneration">Generation number. 1 at enrollment.</param>
    public static UserSecretBundle Create(DateTimeOffset createdAt, uint keyGeneration = 1)
    {
        ArgumentOutOfRangeException.ThrowIfLessThan(keyGeneration, 1u);

        var parameters = new KeyCreationParameters
        {
            ExportPolicy = KeyExportPolicies.AllowPlaintextExport,
        };

        var encryption = Key.Create(KeyAgreementAlgorithm.X25519, parameters);

        Key signing;
        try
        {
            signing = Key.Create(SignatureAlgorithm.Ed25519, parameters);
        }
        catch
        {
            encryption.Dispose();
            throw;
        }

        return new UserSecretBundle(
            encryption,
            signing,
            keyGeneration,
            KeyLogChain.TruncateTimestamp(createdAt));
    }

    /// <summary>
    /// Derives the key that encrypts this client's on-disk cache. See docs/crypto.md §3.2.
    /// </summary>
    /// <remarks>
    /// <para>
    /// Derived from the bundle rather than from the passphrase master key, because the bundle is the one
    /// thing every unlock path ends up holding. A device unlock opens a <see cref="SealTo"/> wrap and never
    /// sees a master key at all, so under the old derivation it could open the identity and still not read
    /// the cache it had itself written; a recovery-code unlock would derive a different key from a different
    /// master and silently orphan every cached row. Keying on the bundle makes the cache readable through
    /// whichever door was used, and unreadable exactly when the identity itself has been rotated — which is
    /// the correct time to discard it.
    /// </para>
    /// <para>
    /// <b>Extract-and-expand, not expand alone.</b> <see cref="MasterKey"/> can expand directly because an
    /// Argon2id output is already uniformly random over its whole length. The bundle's encoding is not: it
    /// opens with a fixed 14-byte label and carries a version, a generation and a timestamp before it
    /// reaches any key material. HKDF's extract step is what turns that into a pseudorandom key, and
    /// skipping it would be feeding structured input to a function that assumes it has none.
    /// </para>
    /// </remarks>
    public void DeriveLocalCacheKey(Span<byte> destination)
    {
        Alive();

        if (destination.Length != CryptoSpec.SymmetricKeySize)
        {
            throw new ArgumentException(
                $"The cache key is {CryptoSpec.SymmetricKeySize} bytes, got {destination.Length}.",
                nameof(destination));
        }

        Span<byte> encoded = stackalloc byte[EncodedLength];
        try
        {
            Encode(encoded);

            HKDF.DeriveKey(
                HashAlgorithmName.SHA512,
                encoded,
                destination,
                salt: default,
                info: CryptoSpec.DerivationLabels.LocalCache);
        }
        finally
        {
            CryptographicOperations.ZeroMemory(encoded);
        }
    }

    /// <summary>
    /// Wraps the bundle under a symmetric key-encryption key.
    /// </summary>
    /// <remarks>
    /// Used for the passphrase and recovery-code wraps, whose keys are derived rather than sealed
    /// to. The plaintext exists only inside this method and is zeroed before it returns.
    /// </remarks>
    public byte[] WrapUnder(ReadOnlySpan<byte> keyEncryptionKey, in AadDescriptor descriptor)
    {
        Alive();

        Span<byte> plaintext = stackalloc byte[EncodedLength];
        try
        {
            Encode(plaintext);
            return DshCrypto.Seal(keyEncryptionKey, plaintext, descriptor);
        }
        finally
        {
            CryptographicOperations.ZeroMemory(plaintext);
        }
    }

    /// <summary>
    /// Seals the bundle to a public key.
    /// </summary>
    /// <remarks>
    /// Used for device and escrow wraps. Anonymous-sender by construction, which is fine here: the
    /// recipient is the same user, so there is nothing to attribute.
    /// </remarks>
    public byte[] SealTo(ReadOnlySpan<byte> recipientPublicKey, in AadDescriptor descriptor)
    {
        Alive();

        Span<byte> plaintext = stackalloc byte[EncodedLength];
        try
        {
            Encode(plaintext);
            return DshCrypto.SealTo(recipientPublicKey, plaintext, descriptor);
        }
        finally
        {
            CryptographicOperations.ZeroMemory(plaintext);
        }
    }

    /// <summary>
    /// Opens a wrap made by <see cref="WrapUnder"/>.
    /// </summary>
    /// <returns>
    /// The bundle, or <see langword="null"/> if the key is wrong, the descriptor does not match, or
    /// the encoding is unrecognised. Never throws on bad input: the wrap arrives from a server that
    /// is explicitly not trusted, so a failed tag is an expected outcome — most often simply a
    /// mistyped passphrase.
    /// </returns>
    public static UserSecretBundle? TryOpenUnder(
        ReadOnlySpan<byte> keyEncryptionKey,
        ReadOnlySpan<byte> envelope,
        in AadDescriptor descriptor)
    {
        var plaintext = DshCrypto.Open(keyEncryptionKey, envelope, descriptor);
        return FromPlaintext(plaintext);
    }

    /// <summary>Opens a wrap made by <see cref="SealTo"/>, using the recipient's private key.</summary>
    public static UserSecretBundle? TryOpenSealed(
        Key recipientKey,
        ReadOnlySpan<byte> envelope,
        in AadDescriptor descriptor)
    {
        var plaintext = DshCrypto.OpenSealed(recipientKey, envelope, descriptor);
        return FromPlaintext(plaintext);
    }

    /// <inheritdoc />
    public void Dispose()
    {
        if (disposed)
        {
            return;
        }

        disposed = true;
        encryptionKey.Dispose();
        signingKey.Dispose();
    }

    private void Encode(Span<byte> destination)
    {
        Label.CopyTo(destination);
        BinaryPrimitives.WriteUInt16BigEndian(destination[OffsetVersion..], CurrentVersion);
        BinaryPrimitives.WriteUInt32BigEndian(destination[OffsetKeyGeneration..], KeyGeneration);
        BinaryPrimitives.WriteInt64BigEndian(destination[OffsetCreatedAt..], CreatedAt.ToUnixTimeMilliseconds());

        // Ed25519's raw private key is the 32-byte seed, and X25519's is the 32-byte scalar, so both
        // fit the fixed layout. Exported straight into the caller's buffer, which is zeroed there.
        ExportInto(encryptionKey, destination.Slice(OffsetEncryptionKey, CryptoSpec.SymmetricKeySize));
        ExportInto(signingKey, destination.Slice(OffsetSigningKey, CryptoSpec.SymmetricKeySize));
    }

    private static void ExportInto(Key key, Span<byte> destination)
    {
        var exported = key.Export(KeyBlobFormat.RawPrivateKey);
        try
        {
            if (exported.Length != destination.Length)
            {
                throw new CryptographicException(
                    $"Expected a {destination.Length}-byte raw private key, got {exported.Length}.");
            }

            exported.CopyTo(destination);
        }
        finally
        {
            CryptographicOperations.ZeroMemory(exported);
        }
    }

    private static UserSecretBundle? FromPlaintext(byte[]? plaintext)
    {
        if (plaintext is null)
        {
            return null;
        }

        try
        {
            return TryDecode(plaintext);
        }
        finally
        {
            CryptographicOperations.ZeroMemory(plaintext);
        }
    }

    /// <remarks>
    /// Fails closed on anything unexpected. A bundle is the root of everything a user can read, so
    /// guessing at a malformed one is never the right move.
    /// </remarks>
    private static bool IsWellFormed(ReadOnlySpan<byte> encoded) =>
        encoded.Length == EncodedLength
        && encoded[..Label.Length].SequenceEqual(Label)
        && BinaryPrimitives.ReadUInt16BigEndian(encoded[OffsetVersion..]) == CurrentVersion
        && BinaryPrimitives.ReadUInt32BigEndian(encoded[OffsetKeyGeneration..]) >= 1;

    /// <remarks>
    /// Internal rather than private so the golden-vector generator can build a bundle from fixed bytes.
    /// The cache key derives from this encoding, and a vector for it is worth nothing unless the input is
    /// pinned — <see cref="Create"/> draws fresh randomness, so it cannot produce a reproducible one.
    /// </remarks>
    internal static UserSecretBundle? TryDecode(ReadOnlySpan<byte> encoded)
    {
        if (!IsWellFormed(encoded))
        {
            return null;
        }

        var keyGeneration = BinaryPrimitives.ReadUInt32BigEndian(encoded[OffsetKeyGeneration..]);

        var createdAt = DateTimeOffset.FromUnixTimeMilliseconds(
            BinaryPrimitives.ReadInt64BigEndian(encoded[OffsetCreatedAt..]));

        var parameters = new KeyCreationParameters
        {
            ExportPolicy = KeyExportPolicies.AllowPlaintextExport,
        };

        Key? encryption = null;
        Key? signing = null;

        try
        {
            encryption = Key.Import(
                KeyAgreementAlgorithm.X25519,
                encoded.Slice(OffsetEncryptionKey, CryptoSpec.SymmetricKeySize),
                KeyBlobFormat.RawPrivateKey,
                parameters);

            signing = Key.Import(
                SignatureAlgorithm.Ed25519,
                encoded.Slice(OffsetSigningKey, CryptoSpec.SymmetricKeySize),
                KeyBlobFormat.RawPrivateKey,
                parameters);

            var bundle = new UserSecretBundle(encryption, signing, keyGeneration, createdAt);

            // Ownership transferred; do not dispose in the catch below.
            encryption = null;
            signing = null;

            return bundle;
        }
        catch (FormatException)
        {
            return null;
        }
        finally
        {
            encryption?.Dispose();
            signing?.Dispose();
        }
    }

    private UserSecretBundle Alive()
    {
        ObjectDisposedException.ThrowIf(disposed, this);
        return this;
    }
}