namespace DodoSSH.Crypto;

/// <summary>
/// Framing of the DSH1 envelope. See docs/crypto.md §5, which is normative.
/// </summary>
/// <remarks>
/// <para>
/// Layout, all integers big-endian:
/// </para>
/// <code>
///  offset size  field
///       0    4  magic          "DSH1"
///       4    1  alg_id
///       5    1  flags          reserved, must be zero
///    [alg_id = 3 only]
///       6   32  ephemeral_pk
///       —    —  nonce          24 bytes for alg 1 and 3, 12 for alg 2
///       —    n  ciphertext     including the trailing 16-byte tag
/// </code>
/// <para>
/// This type does framing only; it performs no cryptography. Readers reject unknown
/// algorithms and any non-zero flag bit, so an envelope that is not fully understood fails
/// closed rather than being partially interpreted.
/// </para>
/// </remarks>
public static class DshEnvelope
{
    /// <summary>Bytes before the algorithm-specific portion.</summary>
    public const int HeaderSize = 6;

    private const int OffsetAlgorithm = 4;
    private const int OffsetFlags = 5;
    private const int OffsetEphemeralPublicKey = 6;

    /// <summary>Nonce length for XChaCha20-Poly1305 and SealTo.</summary>
    public const int XChaChaNonceSize = 24;

    /// <summary>Nonce length for AES-256-GCM.</summary>
    public const int AesGcmNonceSize = 12;

    /// <summary>Returns the nonce length used by an algorithm.</summary>
    public static int NonceSizeFor(CryptoSpec.AlgorithmId algorithm) => algorithm switch
    {
        CryptoSpec.AlgorithmId.XChaCha20Poly1305 => XChaChaNonceSize,
        CryptoSpec.AlgorithmId.SealToX25519 => XChaChaNonceSize,
        CryptoSpec.AlgorithmId.Aes256Gcm => AesGcmNonceSize,
        _ => throw new ArgumentOutOfRangeException(nameof(algorithm), algorithm, "Unknown algorithm."),
    };

    /// <summary>True when the algorithm carries an ephemeral public key in its header.</summary>
    public static bool HasEphemeralPublicKey(CryptoSpec.AlgorithmId algorithm) =>
        algorithm == CryptoSpec.AlgorithmId.SealToX25519;

    /// <summary>Total prefix length before the ciphertext for a given algorithm.</summary>
    public static int PrefixSizeFor(CryptoSpec.AlgorithmId algorithm) =>
        HeaderSize
        + (HasEphemeralPublicKey(algorithm) ? CryptoSpec.PublicKeySize : 0)
        + NonceSizeFor(algorithm);

    /// <summary>Smallest legal envelope for an algorithm: prefix plus a bare tag.</summary>
    public static int MinimumSizeFor(CryptoSpec.AlgorithmId algorithm) =>
        PrefixSizeFor(algorithm) + CryptoSpec.TagSize;

    /// <summary>Exact size of an envelope for an algorithm and ciphertext length.</summary>
    public static int SizeFor(CryptoSpec.AlgorithmId algorithm, int ciphertextLength) =>
        PrefixSizeFor(algorithm) + ciphertextLength;

    /// <summary>
    /// Writes an envelope.
    /// </summary>
    /// <param name="destination">Buffer receiving the envelope.</param>
    /// <param name="algorithm">Construction used.</param>
    /// <param name="nonce">Nonce, whose length must match the algorithm.</param>
    /// <param name="ciphertext">Ciphertext including its trailing tag.</param>
    /// <param name="ephemeralPublicKey">
    /// Ephemeral X25519 public key, required for <see cref="CryptoSpec.AlgorithmId.SealToX25519"/>
    /// and forbidden otherwise.
    /// </param>
    /// <returns>Number of bytes written.</returns>
    public static int Write(
        Span<byte> destination,
        CryptoSpec.AlgorithmId algorithm,
        ReadOnlySpan<byte> nonce,
        ReadOnlySpan<byte> ciphertext,
        ReadOnlySpan<byte> ephemeralPublicKey = default)
    {
        var expectedNonce = NonceSizeFor(algorithm);
        if (nonce.Length != expectedNonce)
        {
            throw new ArgumentException(
                $"{algorithm} requires a {expectedNonce}-byte nonce, got {nonce.Length}.",
                nameof(nonce));
        }

        var wantsEphemeral = HasEphemeralPublicKey(algorithm);
        if (wantsEphemeral && ephemeralPublicKey.Length != CryptoSpec.PublicKeySize)
        {
            throw new ArgumentException(
                $"{algorithm} requires a {CryptoSpec.PublicKeySize}-byte ephemeral public key.",
                nameof(ephemeralPublicKey));
        }

        if (!wantsEphemeral && !ephemeralPublicKey.IsEmpty)
        {
            throw new ArgumentException(
                $"{algorithm} does not carry an ephemeral public key.",
                nameof(ephemeralPublicKey));
        }

        if (ciphertext.Length < CryptoSpec.TagSize)
        {
            throw new ArgumentException(
                $"Ciphertext must include a {CryptoSpec.TagSize}-byte tag.",
                nameof(ciphertext));
        }

        var total = SizeFor(algorithm, ciphertext.Length);
        if (destination.Length < total)
        {
            throw new ArgumentException($"Destination must be at least {total} bytes.", nameof(destination));
        }

        CryptoSpec.EnvelopeMagic.CopyTo(destination);
        destination[OffsetAlgorithm] = (byte)algorithm;
        destination[OffsetFlags] = 0;

        var cursor = HeaderSize;
        if (wantsEphemeral)
        {
            ephemeralPublicKey.CopyTo(destination[cursor..]);
            cursor += CryptoSpec.PublicKeySize;
        }

        nonce.CopyTo(destination[cursor..]);
        cursor += nonce.Length;

        ciphertext.CopyTo(destination[cursor..]);
        return cursor + ciphertext.Length;
    }

    /// <summary>Writes an envelope into a new array.</summary>
    public static byte[] Write(
        CryptoSpec.AlgorithmId algorithm,
        ReadOnlySpan<byte> nonce,
        ReadOnlySpan<byte> ciphertext,
        ReadOnlySpan<byte> ephemeralPublicKey = default)
    {
        var buffer = new byte[SizeFor(algorithm, ciphertext.Length)];
        Write(buffer, algorithm, nonce, ciphertext, ephemeralPublicKey);
        return buffer;
    }

    /// <summary>
    /// Parses an envelope without copying. Returns false for anything malformed or not fully
    /// understood, rather than throwing, because envelopes arrive from an untrusted server.
    /// </summary>
    public static bool TryRead(ReadOnlySpan<byte> envelope, out DshEnvelopeView view)
    {
        view = default;

        if (envelope.Length < HeaderSize)
        {
            return false;
        }

        if (!envelope[..CryptoSpec.EnvelopeMagic.Length].SequenceEqual(CryptoSpec.EnvelopeMagic))
        {
            return false;
        }

        // Reject unknown flag bits: an envelope we do not fully understand must fail closed.
        if (envelope[OffsetFlags] != 0)
        {
            return false;
        }

        var algorithmByte = envelope[OffsetAlgorithm];
        if (!Enum.IsDefined(typeof(CryptoSpec.AlgorithmId), algorithmByte))
        {
            return false;
        }

        var algorithm = (CryptoSpec.AlgorithmId)algorithmByte;
        if (algorithm == CryptoSpec.AlgorithmId.Unspecified)
        {
            return false;
        }

        if (envelope.Length < MinimumSizeFor(algorithm))
        {
            return false;
        }

        var cursor = HeaderSize;
        var ephemeral = ReadOnlySpan<byte>.Empty;
        if (HasEphemeralPublicKey(algorithm))
        {
            ephemeral = envelope.Slice(cursor, CryptoSpec.PublicKeySize);
            cursor += CryptoSpec.PublicKeySize;
        }

        var nonceSize = NonceSizeFor(algorithm);
        var nonce = envelope.Slice(cursor, nonceSize);
        cursor += nonceSize;

        view = new DshEnvelopeView(algorithm, ephemeral, nonce, envelope[cursor..]);
        return true;
    }
}
