using System.Security.Cryptography;
using NSec.Cryptography;

namespace DodoSSH.Crypto;

/// <summary>
/// The DSH1 operations: symmetric content encryption and anonymous-sender key wrapping.
/// </summary>
/// <remarks>
/// <para>
/// docs/crypto.md is normative. Every operation takes an <see cref="AadDescriptor"/> rather
/// than a raw AAD, so a caller cannot omit the binding that stops a server relocating
/// ciphertext between rows or key generations.
/// </para>
/// <para>
/// Decryption returns <see langword="null"/> on authentication failure rather than throwing.
/// Ciphertext arrives from a server that is explicitly not trusted, so a failed tag is an
/// expected outcome to be handled, not an exceptional one.
/// </para>
/// </remarks>
public static class DshCrypto
{
    private static AeadAlgorithm Aead => AeadAlgorithm.XChaCha20Poly1305;

    private static KeyAgreementAlgorithm Agreement => KeyAgreementAlgorithm.X25519;

    /// <summary>
    /// Encrypts under a known symmetric key, producing a complete DSH1 envelope
    /// (<see cref="CryptoSpec.AlgorithmId.XChaCha20Poly1305"/>).
    /// </summary>
    /// <param name="key">32-byte content key.</param>
    /// <param name="plaintext">Data to protect.</param>
    /// <param name="descriptor">Identifies the row this ciphertext belongs to.</param>
    public static byte[] Seal(ReadOnlySpan<byte> key, ReadOnlySpan<byte> plaintext, in AadDescriptor descriptor)
    {
        RequireSymmetricKey(key);

        Span<byte> aad = stackalloc byte[CryptoSpec.DigestSize];
        descriptor.ComputeAad(aad);

        // A 192-bit nonce is why a random nonce per message is safe with no counter to track.
        Span<byte> nonce = stackalloc byte[DshEnvelope.XChaChaNonceSize];
        RandomNumberGenerator.Fill(nonce);

        using var aeadKey = Key.Import(Aead, key, KeyBlobFormat.RawSymmetricKey);
        var ciphertext = Aead.Encrypt(aeadKey, nonce, aad, plaintext);

        return DshEnvelope.Write(CryptoSpec.AlgorithmId.XChaCha20Poly1305, nonce, ciphertext);
    }

    /// <summary>
    /// Opens an envelope produced by <see cref="Seal"/>.
    /// </summary>
    /// <returns>The plaintext, or <see langword="null"/> if the envelope is malformed, uses
    /// another construction, or fails authentication under this descriptor.</returns>
    public static byte[]? Open(ReadOnlySpan<byte> key, ReadOnlySpan<byte> envelope, in AadDescriptor descriptor)
    {
        RequireSymmetricKey(key);

        if (!DshEnvelope.TryRead(envelope, out var view))
        {
            return null;
        }

        if (view.Algorithm != CryptoSpec.AlgorithmId.XChaCha20Poly1305)
        {
            return null;
        }

        Span<byte> aad = stackalloc byte[CryptoSpec.DigestSize];
        descriptor.ComputeAad(aad);

        using var aeadKey = Key.Import(Aead, key, KeyBlobFormat.RawSymmetricKey);
        return Aead.Decrypt(aeadKey, view.Nonce, aad, view.Ciphertext);
    }

    /// <summary>
    /// Seals to a recipient's X25519 public key, producing a
    /// <see cref="CryptoSpec.AlgorithmId.SealToX25519"/> envelope. See docs/crypto.md §6.
    /// </summary>
    /// <remarks>
    /// Anonymous-sender by construction: this proves nothing about who created the envelope.
    /// Callers that need attribution — grants, in particular — must additionally attach a
    /// detached Ed25519 signature, or a server could fabricate a grant undetectably.
    /// </remarks>
    public static byte[] SealTo(
        ReadOnlySpan<byte> recipientPublicKey,
        ReadOnlySpan<byte> plaintext,
        in AadDescriptor descriptor)
    {
        RequirePublicKey(recipientPublicKey);

        Span<byte> aad = stackalloc byte[CryptoSpec.DigestSize];
        descriptor.ComputeAad(aad);

        var recipient = PublicKey.Import(Agreement, recipientPublicKey, KeyBlobFormat.RawPublicKey);

        // The ephemeral key must be exportable: its public half goes into the envelope.
        using var ephemeral = Key.Create(
            Agreement,
            new KeyCreationParameters { ExportPolicy = KeyExportPolicies.AllowPlaintextExport });

        var ephemeralPublic = ephemeral.PublicKey.Export(KeyBlobFormat.RawPublicKey);

        using var shared = Agreement.Agree(ephemeral, recipient)
            ?? throw new CryptographicException("X25519 agreement failed; the recipient key is invalid.");

        Span<byte> derived = stackalloc byte[CryptoSpec.SymmetricKeySize];
        DeriveSealKey(shared, ephemeralPublic, recipientPublicKey, aad, derived);

        Span<byte> nonce = stackalloc byte[DshEnvelope.XChaChaNonceSize];
        RandomNumberGenerator.Fill(nonce);

        using var contentKey = Key.Import(Aead, derived, KeyBlobFormat.RawSymmetricKey);
        var ciphertext = Aead.Encrypt(contentKey, nonce, aad, plaintext);

        CryptographicOperations.ZeroMemory(derived);

        return DshEnvelope.Write(
            CryptoSpec.AlgorithmId.SealToX25519,
            nonce,
            ciphertext,
            ephemeralPublic);
    }

    /// <summary>
    /// Opens an envelope produced by <see cref="SealTo"/> using the recipient's private key.
    /// </summary>
    /// <returns>The plaintext, or <see langword="null"/> on any failure.</returns>
    public static byte[]? OpenSealed(Key recipientKey, ReadOnlySpan<byte> envelope, in AadDescriptor descriptor)
    {
        ArgumentNullException.ThrowIfNull(recipientKey);

        if (!DshEnvelope.TryRead(envelope, out var view))
        {
            return null;
        }

        if (view.Algorithm != CryptoSpec.AlgorithmId.SealToX25519)
        {
            return null;
        }

        Span<byte> aad = stackalloc byte[CryptoSpec.DigestSize];
        descriptor.ComputeAad(aad);

        PublicKey ephemeralPublic;
        try
        {
            ephemeralPublic = PublicKey.Import(Agreement, view.EphemeralPublicKey, KeyBlobFormat.RawPublicKey);
        }
        catch (FormatException)
        {
            return null;
        }

        using var shared = Agreement.Agree(recipientKey, ephemeralPublic);
        if (shared is null)
        {
            // All-zero agreement: a low-order or otherwise invalid ephemeral key.
            return null;
        }

        var recipientPublic = recipientKey.PublicKey.Export(KeyBlobFormat.RawPublicKey);

        Span<byte> derived = stackalloc byte[CryptoSpec.SymmetricKeySize];
        DeriveSealKey(shared, view.EphemeralPublicKey, recipientPublic, aad, derived);

        using var contentKey = Key.Import(Aead, derived, KeyBlobFormat.RawSymmetricKey);
        var plaintext = Aead.Decrypt(contentKey, view.Nonce, aad, view.Ciphertext);

        CryptographicOperations.ZeroMemory(derived);
        return plaintext;
    }

    /// <summary>
    /// Computes an identity fingerprint over a user's two public keys. See docs/crypto.md §8.
    /// </summary>
    public static byte[] ComputeFingerprint(
        ReadOnlySpan<byte> x25519PublicKey,
        ReadOnlySpan<byte> ed25519PublicKey)
    {
        RequirePublicKey(x25519PublicKey);
        RequirePublicKey(ed25519PublicKey);

        var label = CryptoSpec.DerivationLabels.Fingerprint;

        Span<byte> input = stackalloc byte[label.Length + (CryptoSpec.PublicKeySize * 2)];
        label.CopyTo(input);
        x25519PublicKey.CopyTo(input[label.Length..]);
        ed25519PublicKey.CopyTo(input[(label.Length + CryptoSpec.PublicKeySize)..]);

        var fingerprint = new byte[CryptoSpec.DigestSize];
        SHA256.HashData(input, fingerprint);
        return fingerprint;
    }

    /// <summary>
    /// HKDF-SHA256 over the X25519 shared secret, salted with both public keys and bound to
    /// the AAD. Specified rather than reusing libsodium's sealed box, whose KDF covers only the
    /// key pair and would not carry the AAD binding.
    /// </summary>
    private static void DeriveSealKey(
        SharedSecret shared,
        ReadOnlySpan<byte> ephemeralPublicKey,
        ReadOnlySpan<byte> recipientPublicKey,
        ReadOnlySpan<byte> aad,
        Span<byte> destination)
    {
        Span<byte> salt = stackalloc byte[CryptoSpec.PublicKeySize * 2];
        ephemeralPublicKey.CopyTo(salt);
        recipientPublicKey.CopyTo(salt[CryptoSpec.PublicKeySize..]);

        var prefix = CryptoSpec.DerivationLabels.SealTo;
        Span<byte> info = stackalloc byte[prefix.Length + aad.Length];
        prefix.CopyTo(info);
        aad.CopyTo(info[prefix.Length..]);

        KeyDerivationAlgorithm.HkdfSha256.DeriveBytes(shared, salt, info, destination);
    }

    private static void RequireSymmetricKey(ReadOnlySpan<byte> key)
    {
        if (key.Length != CryptoSpec.SymmetricKeySize)
        {
            throw new ArgumentException(
                $"Key must be {CryptoSpec.SymmetricKeySize} bytes, got {key.Length}.",
                nameof(key));
        }
    }

    private static void RequirePublicKey(ReadOnlySpan<byte> publicKey)
    {
        if (publicKey.Length != CryptoSpec.PublicKeySize)
        {
            throw new ArgumentException(
                $"Public key must be {CryptoSpec.PublicKeySize} bytes, got {publicKey.Length}.",
                nameof(publicKey));
        }
    }
}