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

namespace DodoSSH.Crypto;

/// <summary>
/// The canonical encoding of an identity key statement, and the identity-provider binding value
/// derived from it. See docs/crypto.md §7.1.
/// </summary>
/// <remarks>
/// <para>
/// This encoding is the pivot of the whole public-key trust story. The client hashes it and uses
/// the result as the <c>nonce</c> of a fresh OIDC authorization, so the identity provider ends up
/// signing an assertion over exactly these public keys. The DodoSSH server cannot mint identity
/// provider signatures, so it cannot fabricate a key for a user who never enrolled.
/// </para>
/// <para>
/// It must therefore be reproducible byte-for-byte by every implementation. JSON is unsuitable:
/// property order, number formatting, Unicode escaping and whitespace all vary between
/// serialisers, and two implementations that disagree by one byte produce two different nonces
/// and an enrollment that can never be verified. Hence a fixed binary encoding with explicit
/// length prefixes, where no field value can forge a field boundary.
/// </para>
/// </remarks>
public static class KeyStatementCodec
{
    /// <summary>Domain-separating prefix of the canonical encoding.</summary>
    public static ReadOnlySpan<byte> Label => "dsh1/keystatement/v1"u8;

    /// <summary>Highest statement version this implementation encodes.</summary>
    public const int CurrentVersion = 1;

    /// <summary>Length of the binding value: a SHA-256 digest.</summary>
    public const int BindingLength = CryptoSpec.DigestSize;

    /// <summary>Length of the base64url binding, as it appears in the <c>nonce</c> claim.</summary>
    public const int NonceLength = 43;

    /// <summary>version (u16) + keyGeneration (u32) + createdAt (i64) + two public keys.</summary>
    private const int FixedBlockLength =
        sizeof(ushort) + sizeof(uint) + sizeof(long) + (CryptoSpec.PublicKeySize * 2);

    /// <summary>Presence byte plus a u32 length prefix.</summary>
    private const int StringHeaderLength = 1 + sizeof(uint);

    /// <summary>
    /// Writes the canonical encoding.
    /// </summary>
    /// <remarks>
    /// <para>
    /// Layout, all integers big-endian:
    /// <c>"dsh1/keystatement/v1"</c> | u16 version | u32 keyGeneration |
    /// i64 createdAt (Unix milliseconds, UTC) | x25519 pk (32) | ed25519 pk (32) |
    /// issuer | subject | email | deviceName.
    /// </para>
    /// <para>
    /// Each string is a presence byte — 0 for absent, 1 for present — followed, when present, by a
    /// u32 length and that many UTF-8 bytes. The presence byte is what distinguishes an absent
    /// email from an empty one; without it the encoding would not be injective, and two different
    /// statements could share a binding.
    /// </para>
    /// <para>
    /// <see cref="DateTimeOffset.ToUnixTimeMilliseconds"/> normalises to UTC and truncates, so the
    /// offset a client happens to hold and any sub-millisecond precision are both irrelevant to
    /// the result. That matters because PostgreSQL stores microseconds: a value that survived a
    /// database round trip must still hash identically.
    /// </para>
    /// </remarks>
    /// <param name="statement">The statement to encode.</param>
    /// <exception cref="ArgumentOutOfRangeException">A version, generation or key length is invalid.</exception>
    public static byte[] Encode(KeyStatementFields statement)
    {
        ArgumentNullException.ThrowIfNull(statement);
        Validate(statement);

        var issuer = Encoding.UTF8.GetBytes(statement.Issuer);
        var subject = Encoding.UTF8.GetBytes(statement.Subject);
        var email = statement.Email is null ? null : Encoding.UTF8.GetBytes(statement.Email);
        var deviceName = Encoding.UTF8.GetBytes(statement.DeviceName);

        var length = Label.Length
            + FixedBlockLength
            + StringBlockLength(issuer)
            + StringBlockLength(subject)
            + StringBlockLength(email)
            + StringBlockLength(deviceName);

        var buffer = new byte[length];
        var span = buffer.AsSpan();

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

        BinaryPrimitives.WriteUInt16BigEndian(span[offset..], (ushort)statement.Version);
        offset += sizeof(ushort);

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

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

        statement.EncryptionPublicKey.CopyTo(span[offset..]);
        offset += CryptoSpec.PublicKeySize;

        statement.SigningPublicKey.CopyTo(span[offset..]);
        offset += CryptoSpec.PublicKeySize;

        offset = WriteString(span, offset, issuer);
        offset = WriteString(span, offset, subject);
        offset = WriteString(span, offset, email);
        offset = WriteString(span, offset, deviceName);

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

        return buffer;
    }

    /// <summary>
    /// Computes the identity-provider binding: SHA-256 over the canonical encoding.
    /// </summary>
    public static byte[] ComputeBinding(KeyStatementFields statement) =>
        SHA256.HashData(Encode(statement));

    /// <summary>
    /// Computes the identity-provider binding from an encoding already produced by
    /// <see cref="Encode"/>.
    /// </summary>
    /// <remarks>
    /// For callers that also need the encoding itself — to verify a signature over it, say — so
    /// that the hash algorithm stays a decision of this type rather than being restated at the call
    /// site.
    /// </remarks>
    public static byte[] ComputeBinding(ReadOnlySpan<byte> canonicalEncoding) =>
        SHA256.HashData(canonicalEncoding);

    /// <summary>
    /// Renders a binding as the string that must appear in the OIDC <c>nonce</c> claim.
    /// </summary>
    /// <remarks>
    /// Base64url without padding, because the nonce travels in an authorization request query
    /// string where <c>+</c>, <c>/</c> and <c>=</c> all need escaping and some providers mangle
    /// them.
    /// </remarks>
    public static string ToNonce(ReadOnlySpan<byte> binding)
    {
        if (binding.Length != BindingLength)
        {
            throw new ArgumentException(
                $"A binding is {BindingLength} bytes, got {binding.Length}.",
                nameof(binding));
        }

        return Base64Url.EncodeToString(binding);
    }

    /// <summary>Computes the expected <c>nonce</c> claim value for a statement.</summary>
    public static string ComputeNonce(KeyStatementFields statement) =>
        ToNonce(ComputeBinding(statement));

    private static void Validate(KeyStatementFields statement)
    {
        if (statement.Version is < 1 or > ushort.MaxValue)
        {
            throw new ArgumentOutOfRangeException(
                nameof(statement),
                statement.Version,
                $"Statement version must be between 1 and {ushort.MaxValue}.");
        }

        if (statement.KeyGeneration < 1)
        {
            throw new ArgumentOutOfRangeException(
                nameof(statement),
                statement.KeyGeneration,
                "Key generation must be at least 1.");
        }

        if (statement.Issuer is null || statement.Subject is null || statement.DeviceName is null)
        {
            throw new ArgumentException(
                "Issuer, subject and device name are all required.",
                nameof(statement));
        }

        RequirePublicKey(statement.EncryptionPublicKey, nameof(statement));
        RequirePublicKey(statement.SigningPublicKey, nameof(statement));
    }

    private static int StringBlockLength(byte[]? value) =>
        value is null ? 1 : StringHeaderLength + value.Length;

    private static int WriteString(Span<byte> destination, int offset, byte[]? value)
    {
        if (value is null)
        {
            destination[offset] = 0;
            return offset + 1;
        }

        destination[offset] = 1;
        BinaryPrimitives.WriteUInt32BigEndian(destination[(offset + 1)..], (uint)value.Length);
        value.CopyTo(destination[(offset + StringHeaderLength)..]);

        return offset + StringHeaderLength + value.Length;
    }

    private static void RequirePublicKey(byte[]? publicKey, string parameterName)
    {
        if (publicKey is null || publicKey.Length != CryptoSpec.PublicKeySize)
        {
            throw new ArgumentOutOfRangeException(
                parameterName,
                publicKey?.Length ?? 0,
                $"Public keys must be {CryptoSpec.PublicKeySize} bytes.");
        }
    }
}
