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

namespace DodoSSH.Crypto;

/// <summary>
/// Why a vault key grant exists.
/// </summary>
/// <remarks>
/// <para>
/// Named <c>GrantPurpose</c> rather than <c>GrantKind</c> only to avoid colliding with
/// <c>DodoSSH.Domain.GrantKind</c>, which the server uses for the same concept. Both are visible in the
/// server, and an ambiguous name there would need qualifying at every use.
/// </para>
/// <para>
/// The <b>numeric values must match</b> that enum exactly. They are covered by a grant signature, so a
/// renumbering would make every grant of the changed kind fail verification for good. A test pins them.
/// </para>
/// </remarks>
public enum GrantPurpose : byte
{
    /// <summary>Not a legal value.</summary>
    Unspecified = 0,

    /// <summary>Wrapped to a member's identity key.</summary>
    Member = 1,

    /// <summary>Wrapped to a recovery key held by the vault owner.</summary>
    Recovery = 2,

    /// <summary>Wrapped to a team break-glass key.</summary>
    Escrow = 3,
}

/// <summary>
/// The canonical encoding of a vault key grant, as signed by the granter. See docs/crypto.md §7.3.
/// </summary>
/// <remarks>
/// <para>
/// Sealing a vault key is anonymous-sender by construction, so a grant proves nothing about who
/// created it. Without a signature over this tuple a server could fabricate a grant containing a key
/// of its own choosing, and the recipient would unwrap it successfully and be none the wiser. The
/// signature makes that detectable and attributable — it cannot make it impossible, since verifying
/// the contents would require the server to hold the key.
/// </para>
/// <para>
/// The signature covers <c>SHA-256(wrappedKey)</c> rather than the wrapped key itself, so a verifier
/// does not need the vault key to check who issued the grant.
/// </para>
/// </remarks>
public static class GrantStatementCodec
{
    /// <summary>Domain-separating prefix.</summary>
    public static ReadOnlySpan<byte> Label => "dsh1/grant/v1"u8;

    private const int LabelLength = 13;

    /// <summary>Length without a key log head.</summary>
    private const int BaseLength =
        LabelLength
        + sizeof(uint)                  // key generation
        + 1                             // grant kind
        + 16                            // vault id
        + 16                            // grantee user id
        + CryptoSpec.DigestSize         // grantee key fingerprint
        + CryptoSpec.DigestSize         // SHA-256 of the wrapped key
        + 16                            // granter user id
        + CryptoSpec.DigestSize         // granter key fingerprint
        + 1                             // key log head presence
        + sizeof(long);                 // timestamp, Unix milliseconds

    /// <summary>
    /// Writes the canonical encoding.
    /// </summary>
    /// <param name="vaultId">The vault the key belongs to.</param>
    /// <param name="keyGeneration">Generation the grant is for, so a superseded one cannot be replayed.</param>
    /// <param name="kind">Why the grant exists.</param>
    /// <param name="granteeUserId">Who may open it. <see cref="Guid.Empty"/> for a non-member grant.</param>
    /// <param name="granteeKeyFingerprint">The exact identity key it was wrapped to.</param>
    /// <param name="wrappedKey">The sealed vault key; only its digest is covered.</param>
    /// <param name="granterUserId">Who issued it.</param>
    /// <param name="granterKeyFingerprint">The granter's identity key.</param>
    /// <param name="keyLogHead">
    /// The key log head the granter observed, or null. Null for a self-grant, where there is no third
    /// party whose key could have been substituted — and where the log entry that would supply the head
    /// is written in the same transaction, so the client could not have signed over it.
    /// </param>
    /// <param name="grantedAt">Signing time; truncated to milliseconds.</param>
    public static byte[] Encode(
        Guid vaultId,
        uint keyGeneration,
        GrantPurpose kind,
        Guid granteeUserId,
        ReadOnlySpan<byte> granteeKeyFingerprint,
        ReadOnlySpan<byte> wrappedKey,
        Guid granterUserId,
        ReadOnlySpan<byte> granterKeyFingerprint,
        ReadOnlySpan<byte> keyLogHead,
        DateTimeOffset grantedAt)
    {
        Validate(kind, granteeKeyFingerprint, wrappedKey, granterKeyFingerprint, keyLogHead);

        var buffer = new byte[BaseLength + keyLogHead.Length];
        var span = buffer.AsSpan();

        Label.CopyTo(span);
        var offset = LabelLength;

        BinaryPrimitives.WriteUInt32BigEndian(span[offset..], keyGeneration);
        offset += sizeof(uint);

        span[offset++] = (byte)kind;

        offset = WriteGuid(span, offset, vaultId);
        offset = WriteGuid(span, offset, granteeUserId);

        granteeKeyFingerprint.CopyTo(span[offset..]);
        offset += CryptoSpec.DigestSize;

        // The digest, not the key. A verifier must be able to check attribution without holding the
        // vault key.
        SHA256.HashData(wrappedKey, span.Slice(offset, CryptoSpec.DigestSize));
        offset += CryptoSpec.DigestSize;

        offset = WriteGuid(span, offset, granterUserId);

        granterKeyFingerprint.CopyTo(span[offset..]);
        offset += CryptoSpec.DigestSize;

        if (keyLogHead.IsEmpty)
        {
            span[offset++] = 0;
        }
        else
        {
            span[offset++] = 1;
            keyLogHead.CopyTo(span[offset..]);
            offset += CryptoSpec.DigestSize;
        }

        BinaryPrimitives.WriteInt64BigEndian(span[offset..], grantedAt.ToUnixTimeMilliseconds());
        offset += sizeof(long);

        if (offset != buffer.Length)
        {
            throw new InvalidOperationException(
                $"Grant encoding wrote {offset} bytes but reserved {buffer.Length}.");
        }

        return buffer;
    }

    private static void Validate(
        GrantPurpose kind,
        ReadOnlySpan<byte> granteeKeyFingerprint,
        ReadOnlySpan<byte> wrappedKey,
        ReadOnlySpan<byte> granterKeyFingerprint,
        ReadOnlySpan<byte> keyLogHead)
    {
        if (kind == GrantPurpose.Unspecified)
        {
            throw new ArgumentOutOfRangeException(nameof(kind), kind, "A grant kind is required.");
        }

        RequireDigest(granteeKeyFingerprint, nameof(granteeKeyFingerprint));
        RequireDigest(granterKeyFingerprint, nameof(granterKeyFingerprint));

        if (wrappedKey.IsEmpty)
        {
            throw new ArgumentException("A wrapped key is required.", nameof(wrappedKey));
        }

        if (!keyLogHead.IsEmpty && keyLogHead.Length != CryptoSpec.DigestSize)
        {
            throw new ArgumentException(
                $"A key log head is {CryptoSpec.DigestSize} bytes or absent.",
                nameof(keyLogHead));
        }
    }

    /// <summary>Signs a grant encoding.</summary>
    public static byte[] Sign(Key signingKey, ReadOnlySpan<byte> canonicalGrant)
    {
        ArgumentNullException.ThrowIfNull(signingKey);

        return SignatureAlgorithm.Ed25519.Sign(signingKey, BuildMessage(canonicalGrant));
    }

    /// <summary>
    /// Verifies a grant signature against the granter's published signing key.
    /// </summary>
    /// <remarks>
    /// Clients verify these; the server stores them opaquely. Server-side verification would be a
    /// convenience and never the boundary, and would put an asymmetric implementation on a machine
    /// that is supposed to hold no keys.
    /// </remarks>
    public static bool Verify(
        ReadOnlySpan<byte> granterSigningPublicKey,
        ReadOnlySpan<byte> canonicalGrant,
        ReadOnlySpan<byte> signature)
    {
        if (granterSigningPublicKey.Length != CryptoSpec.PublicKeySize
            || signature.Length != CryptoSpec.SignatureSize)
        {
            return false;
        }

        PublicKey publicKey;
        try
        {
            publicKey = PublicKey.Import(
                SignatureAlgorithm.Ed25519,
                granterSigningPublicKey,
                KeyBlobFormat.RawPublicKey);
        }
        catch (FormatException)
        {
            return false;
        }

        return SignatureAlgorithm.Ed25519.Verify(publicKey, BuildMessage(canonicalGrant), signature);
    }

    /// <remarks>
    /// Context-prefixed, so a grant signature can never be replayed as a key statement signature or an
    /// attestation.
    /// </remarks>
    private static byte[] BuildMessage(ReadOnlySpan<byte> canonicalGrant)
    {
        var context = CryptoSpec.SigningContexts.Grant;
        var message = new byte[context.Length + canonicalGrant.Length];

        context.CopyTo(message);
        canonicalGrant.CopyTo(message.AsSpan(context.Length));

        return message;
    }

    private static int WriteGuid(Span<byte> destination, int offset, Guid value)
    {
        // RFC 4122 big-endian, as everywhere else in this specification.
        if (!value.TryWriteBytes(destination[offset..], bigEndian: true, out _))
        {
            throw new InvalidOperationException("Failed to write a grant identifier.");
        }

        return offset + 16;
    }

    private static void RequireDigest(ReadOnlySpan<byte> value, string parameterName)
    {
        if (value.Length != CryptoSpec.DigestSize)
        {
            throw new ArgumentException(
                $"Expected a {CryptoSpec.DigestSize}-byte fingerprint, got {value.Length}.",
                parameterName);
        }
    }
}