using System.Security.Cryptography;

namespace DodoSSH.Crypto;

/// <summary>
/// The master key derived from a vault passphrase, and the subkeys it yields. See docs/crypto.md §3.
/// </summary>
/// <remarks>
/// <para>
/// RAM only. Never persisted, never transmitted, never written to a keystore. What <em>is</em>
/// cached locally is the KDF salt and the wrapped bundle, because unlock has to work with no
/// network: a salt fetched at unlock time would mean an offline launch cannot open the vault, which
/// is the single most common thing a user does on a plane.
/// </para>
/// <para>
/// Argon2id parameters travel with each wrap rather than being global, so raising them later is a
/// per-user migration at next unlock rather than a breaking change — an older client can still open
/// its own wrap. See docs/crypto.md §2 for why the memory unit needs guarding.
/// </para>
/// </remarks>
public sealed class MasterKey : IDisposable
{
    private readonly byte[] material = new byte[CryptoSpec.MasterKeySize];
    private bool disposed;

    private MasterKey(ReadOnlySpan<byte> derived) => derived.CopyTo(material);

    /// <summary>
    /// Stretches a passphrase into the master key.
    /// </summary>
    /// <remarks>
    /// The cost is deliberately high enough to be noticeable — a few hundred milliseconds on a
    /// desktop at the default profile — because this is the only thing standing between a stolen
    /// database dump and every credential in the vault.
    /// </remarks>
    /// <param name="passphrase">The user's vault passphrase.</param>
    /// <param name="salt">Per-wrap salt, at least 16 bytes. Not a secret.</param>
    /// <param name="profile">Argon2id cost parameters.</param>
    public static MasterKey Derive(string passphrase, ReadOnlySpan<byte> salt, Argon2Profile profile)
    {
        ArgumentException.ThrowIfNullOrEmpty(passphrase);
        ArgumentNullException.ThrowIfNull(profile);

        if (salt.Length < CryptoSpec.SaltSize)
        {
            throw new ArgumentException(
                $"Salt must be at least {CryptoSpec.SaltSize} bytes, got {salt.Length}.",
                nameof(salt));
        }

        var derived = profile
            .CreateAlgorithm()
            .DeriveBytes(passphrase, salt, CryptoSpec.MasterKeySize);

        try
        {
            return new MasterKey(derived);
        }
        finally
        {
            CryptographicOperations.ZeroMemory(derived);
        }
    }

    /// <summary>
    /// Derives the key that encrypts the client's on-disk cache.
    /// </summary>
    /// <remarks>
    /// Domain-separated from the bundle's key-encryption key by its HKDF info label, so a cache
    /// record can never be opened with the wrap key or the reverse — the two live in very different
    /// threat models and must not share a key.
    /// </remarks>
    public void DeriveLocalCacheKey(Span<byte> destination) =>
        DeriveSubkey(CryptoSpec.DerivationLabels.LocalCache, destination);

    /// <summary>Wraps a bundle under the passphrase-derived key-encryption key.</summary>
    /// <remarks>
    /// The key-encryption key is never handed out. Callers that only need to wrap or unwrap should
    /// not be holding it, and one that does is one more place it can be leaked or logged.
    /// </remarks>
    public byte[] WrapBundle(UserSecretBundle bundle, in AadDescriptor descriptor)
    {
        ArgumentNullException.ThrowIfNull(bundle);
        ObjectDisposedException.ThrowIf(disposed, this);

        Span<byte> kek = stackalloc byte[CryptoSpec.SymmetricKeySize];
        try
        {
            DeriveSubkey(CryptoSpec.DerivationLabels.PassphraseKek, kek);
            return bundle.WrapUnder(kek, descriptor);
        }
        finally
        {
            CryptographicOperations.ZeroMemory(kek);
        }
    }

    /// <summary>
    /// Opens a bundle wrapped by <see cref="WrapBundle"/>.
    /// </summary>
    /// <returns>
    /// The bundle, or <see langword="null"/> when the passphrase is wrong or the wrap does not
    /// belong to this descriptor. A wrong passphrase is the overwhelmingly common case, so it is a
    /// return value rather than an exception.
    /// </returns>
    public UserSecretBundle? TryOpenBundle(ReadOnlySpan<byte> envelope, in AadDescriptor descriptor)
    {
        ObjectDisposedException.ThrowIf(disposed, this);

        Span<byte> kek = stackalloc byte[CryptoSpec.SymmetricKeySize];
        try
        {
            DeriveSubkey(CryptoSpec.DerivationLabels.PassphraseKek, kek);
            return UserSecretBundle.TryOpenUnder(kek, envelope, descriptor);
        }
        finally
        {
            CryptographicOperations.ZeroMemory(kek);
        }
    }

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

        disposed = true;
        CryptographicOperations.ZeroMemory(material);
    }

    /// <remarks>
    /// <para>
    /// HKDF-Expand rather than the full extract-then-expand: the master key is already a uniformly
    /// random Argon2id output, so there is no low-entropy input left for an extract step to
    /// condition. RFC 5869 §3.3 covers exactly this case.
    /// </para>
    /// <para>
    /// That choice is why the master key is <see cref="CryptoSpec.MasterKeySize"/> bytes rather than
    /// 32: skipping extract means the master key <em>is</em> the PRK, and .NET's
    /// <c>HKDF.Expand</c> rejects a PRK shorter than the hash output.
    /// </para>
    /// </remarks>
    private void DeriveSubkey(ReadOnlySpan<byte> info, Span<byte> destination)
    {
        ObjectDisposedException.ThrowIf(disposed, this);

        if (destination.Length != CryptoSpec.SymmetricKeySize)
        {
            throw new ArgumentException(
                $"Subkeys are {CryptoSpec.SymmetricKeySize} bytes, got {destination.Length}.",
                nameof(destination));
        }

        HKDF.Expand(HashAlgorithmName.SHA512, material, destination, info);
    }
}