using System.Security.Cryptography;

namespace DodoSSH.Crypto;

/// <summary>
/// Per-item data keys and the payloads they protect. See docs/crypto.md §3.
/// </summary>
/// <remarks>
/// <para>
/// Every item version gets its own 32-byte data key, wrapped under the vault key. Four reasons, all
/// of which matter: rotating a vault key re-wraps N × 32 bytes and never touches a content blob, so a
/// 10,000-item vault rotates in a few hundred kilobytes; a single item can be re-wrapped to another
/// key without disturbing the rest; each key encrypts about one message, which makes nonce-collision
/// analysis moot; and a new version gets a new key, so earlier ciphertext stays independently
/// decryptable for history and undo.
/// </para>
/// <para>
/// Be plain about what this does <em>not</em> yet buy. Data keys are wrapped under the vault key, not
/// to individual users, so anyone holding the vault key can decrypt any ciphertext they obtain. Until
/// per-item grants land in M5, an item-level ACL is access control, not cryptographic isolation. The
/// <c>content_key_id</c> column exists from the first migration so that change needs no migration.
/// </para>
/// </remarks>
public static class ItemKeys
{
    /// <summary>Generates a fresh data key. The caller owns and must zero it.</summary>
    public static byte[] CreateDataKey() =>
        RandomNumberGenerator.GetBytes(CryptoSpec.SymmetricKeySize);

    /// <summary>Wraps an item's data key under the vault key.</summary>
    /// <param name="dataKey">The 32-byte data key.</param>
    /// <param name="vaultKey">The vault key it is wrapped under.</param>
    /// <param name="resourceType">What kind of item this is.</param>
    /// <param name="itemId">The item.</param>
    /// <param name="keyGeneration">Vault key generation in force.</param>
    /// <param name="itemVersion">Item version.</param>
    public static byte[] WrapDataKey(
        ReadOnlySpan<byte> dataKey,
        ReadOnlySpan<byte> vaultKey,
        CryptoSpec.AadResourceType resourceType,
        Guid itemId,
        uint keyGeneration,
        uint itemVersion)
    {
        RequireDataKey(dataKey);
        VaultKeys.RequireVaultKey(vaultKey);

        return DshCrypto.Seal(
            vaultKey,
            dataKey,
            DshAad.ItemDataKey(resourceType, itemId, keyGeneration, itemVersion));
    }

    /// <summary>Opens a wrapped data key.</summary>
    /// <returns>The data key, or <see langword="null"/> if it does not belong to this item, version
    /// or generation.</returns>
    public static byte[]? TryUnwrapDataKey(
        ReadOnlySpan<byte> vaultKey,
        ReadOnlySpan<byte> envelope,
        CryptoSpec.AadResourceType resourceType,
        Guid itemId,
        uint keyGeneration,
        uint itemVersion)
    {
        VaultKeys.RequireVaultKey(vaultKey);

        var plaintext = DshCrypto.Open(
            vaultKey,
            envelope,
            DshAad.ItemDataKey(resourceType, itemId, keyGeneration, itemVersion));

        if (plaintext is null)
        {
            return null;
        }

        if (plaintext.Length != CryptoSpec.SymmetricKeySize)
        {
            CryptographicOperations.ZeroMemory(plaintext);
            return null;
        }

        return plaintext;
    }

    /// <summary>Encrypts an item's payload under its data key.</summary>
    /// <param name="dataKey">The item's data key.</param>
    /// <param name="plaintext">The payload.</param>
    /// <param name="resourceType">What kind of item this is.</param>
    /// <param name="itemId">The item.</param>
    /// <param name="dataKeyId">The data key's identifier, stored as <c>content_key_id</c>.</param>
    /// <param name="keyGeneration">Vault key generation in force.</param>
    /// <param name="itemVersion">Item version.</param>
    public static byte[] SealPayload(
        ReadOnlySpan<byte> dataKey,
        ReadOnlySpan<byte> plaintext,
        CryptoSpec.AadResourceType resourceType,
        Guid itemId,
        Guid dataKeyId,
        uint keyGeneration,
        uint itemVersion)
    {
        RequireDataKey(dataKey);

        return DshCrypto.Seal(
            dataKey,
            plaintext,
            DshAad.ItemPayload(resourceType, itemId, dataKeyId, keyGeneration, itemVersion));
    }

    /// <summary>Decrypts an item's payload.</summary>
    /// <returns>The payload, or <see langword="null"/> on any mismatch. A null here after a
    /// successful data key unwrap means the server moved, rolled back or substituted the blob.</returns>
    public static byte[]? TryOpenPayload(
        ReadOnlySpan<byte> dataKey,
        ReadOnlySpan<byte> envelope,
        CryptoSpec.AadResourceType resourceType,
        Guid itemId,
        Guid dataKeyId,
        uint keyGeneration,
        uint itemVersion)
    {
        RequireDataKey(dataKey);

        return DshCrypto.Open(
            dataKey,
            envelope,
            DshAad.ItemPayload(resourceType, itemId, dataKeyId, keyGeneration, itemVersion));
    }

    private static void RequireDataKey(ReadOnlySpan<byte> dataKey)
    {
        if (dataKey.Length != CryptoSpec.SymmetricKeySize)
        {
            throw new ArgumentException(
                $"A data key is {CryptoSpec.SymmetricKeySize} bytes, got {dataKey.Length}.",
                nameof(dataKey));
        }
    }
}
