defmodule AirPlay.V2.Crypto do
  @moduledoc """
  AirPlay 2 cryptographic helpers.

  This module intentionally stays small: HKDF-SHA512, ChaCha20-Poly1305 and AP2
  audio packet encryption. Pairing and RTSP framing build on these primitives.
  """

  @doc "HKDF-SHA512 derive `len` bytes from input keying material."
  @spec hkdf_sha512(binary(), non_neg_integer(), binary(), binary()) :: binary()
  def hkdf_sha512(ikm, len, salt, info) when len > 0 do
    prk = :crypto.mac(:hmac, :sha512, salt, ikm)
    expand_hkdf(prk, info, len, 1, <<>>, <<>>)
  end

  @doc "Compatibility wrapper for HKDF-SHA512 using `(ikm, salt, info, len)` argument order."
  @spec hkdf(binary(), binary(), binary(), non_neg_integer()) :: binary()
  def hkdf(ikm, salt, info, len), do: hkdf_sha512(ikm, len, salt, info)

  @doc "Encrypt with ChaCha20-Poly1305. Returns `{ciphertext, tag}`."
  @spec chacha20_poly1305_encrypt(binary(), binary(), binary(), binary()) :: {binary(), binary()}
  def chacha20_poly1305_encrypt(key, nonce, plaintext, aad \\ <<>>) do
    :crypto.crypto_one_time_aead(:chacha20_poly1305, key, nonce, plaintext, aad, true)
  end

  @doc "Compatibility wrapper for ChaCha20-Poly1305 encryption."
  @spec chacha_encrypt(binary(), binary(), binary(), binary()) :: {binary(), binary()}
  def chacha_encrypt(key, nonce, plaintext, aad \\ <<>>) do
    chacha20_poly1305_encrypt(key, nonce, plaintext, aad)
  end

  @doc "Decrypt with ChaCha20-Poly1305. Returns `{:ok, plaintext}` or `:error`."
  @spec chacha20_poly1305_decrypt(binary(), binary(), binary(), binary(), binary()) ::
          {:ok, binary()} | :error
  def chacha20_poly1305_decrypt(key, nonce, ciphertext, tag, aad \\ <<>>) do
    case :crypto.crypto_one_time_aead(:chacha20_poly1305, key, nonce, ciphertext, aad, tag, false) do
      :error -> :error
      plaintext -> {:ok, plaintext}
    end
  end

  @doc "Compatibility wrapper for ChaCha20-Poly1305 decryption."
  @spec chacha_decrypt(binary(), binary(), binary(), binary(), binary()) :: binary() | :error
  def chacha_decrypt(key, nonce, ciphertext, tag, aad \\ <<>>) do
    case chacha20_poly1305_decrypt(key, nonce, ciphertext, tag, aad) do
      {:ok, plaintext} -> plaintext
      :error -> :error
    end
  end

  @doc "Generate an X25519 keypair as `{public_key, private_key}`."
  @spec x25519_keypair() :: {binary(), binary()}
  def x25519_keypair, do: :crypto.generate_key(:ecdh, :x25519)

  @doc "Compute an X25519 shared secret."
  @spec x25519_shared(binary(), binary()) :: binary()
  def x25519_shared(private_key, public_key),
    do: :crypto.compute_key(:ecdh, public_key, private_key, :x25519)

  @doc "Generate an Ed25519 keypair as `{public_key, private_key}`."
  @spec ed25519_keypair() :: {binary(), binary()}
  def ed25519_keypair, do: :crypto.generate_key(:eddsa, :ed25519)

  @doc "Sign a message with Ed25519."
  @spec ed25519_sign(binary(), binary()) :: binary()
  def ed25519_sign(message, private_key),
    do: :crypto.sign(:eddsa, :none, message, [private_key, :ed25519])

  @doc "Verify an Ed25519 signature."
  @spec ed25519_verify(binary(), binary(), binary()) :: boolean()
  def ed25519_verify(message, signature, public_key),
    do: :crypto.verify(:eddsa, :none, message, signature, [public_key, :ed25519])

  @doc """
  Encrypt an AP2 RTP audio payload.

  The wire format matches owntone and airplay2-rs:

    * nonce = `<<0::32, seq::little-16, 0::48>>`
    * transmitted nonce suffix = `<<seq::little-16, 0::48>>`
    * AAD = `<<rtp_timestamp::32, ssrc::32>>`
    * packet trailer = `tag16 <> nonce8`

  The receiver reconstructs the 12-byte nonce by prefixing four zero bytes to
  the transmitted eight-byte suffix.
  """
  @spec audio_encrypt(binary(), non_neg_integer(), non_neg_integer(), non_neg_integer(), binary()) ::
          binary()
  def audio_encrypt(session_key, rtp_timestamp, ssrc, sequence, payload)
      when byte_size(session_key) == 32 and is_integer(rtp_timestamp) and is_integer(ssrc) and
             is_integer(sequence) do
    nonce_suffix = <<sequence::little-16, 0::48>>
    nonce = <<0::32, nonce_suffix::binary>>
    aad = <<rtp_timestamp::32, ssrc::32>>
    {ciphertext, tag} = chacha20_poly1305_encrypt(session_key, nonce, payload, aad)
    ciphertext <> tag <> nonce_suffix
  end

  @doc "Encrypt an AP2 RTP audio payload using fields parsed from a 12-byte RTP header."
  @spec audio_encrypt(binary(), binary(), binary()) :: binary()
  def audio_encrypt(session_key, rtp_header, payload)
      when byte_size(session_key) == 32 and byte_size(rtp_header) == 12 do
    <<_version, _payload_type, sequence::16, rtp_timestamp::32, ssrc::32>> = rtp_header
    audio_encrypt(session_key, rtp_timestamp, ssrc, sequence, payload)
  end

  def audio_encrypt(session_key, payload, rtp_header)
      when byte_size(session_key) == 32 and byte_size(rtp_header) == 12 do
    audio_encrypt(session_key, rtp_header, payload)
  end

  @doc "Decrypt an AP2 RTP audio payload encrypted by `audio_encrypt/3`."
  @spec audio_decrypt(binary(), binary(), binary()) :: binary() | :error
  def audio_decrypt(session_key, packet, rtp_header)
      when byte_size(session_key) == 32 and byte_size(rtp_header) == 12 and
             byte_size(packet) >= 24 do
    <<_version, _payload_type, _sequence::16, rtp_timestamp::32, ssrc::32>> = rtp_header
    payload_size = byte_size(packet) - 24

    <<ciphertext::binary-size(^payload_size), tag::binary-size(16), nonce_suffix::binary-size(8)>> =
      packet

    nonce = <<0::32, nonce_suffix::binary>>
    aad = <<rtp_timestamp::32, ssrc::32>>
    chacha_decrypt(session_key, nonce, ciphertext, tag, aad)
  end

  defp expand_hkdf(_prk, _info, len, _counter, acc, _previous) when byte_size(acc) >= len do
    binary_part(acc, 0, len)
  end

  defp expand_hkdf(prk, info, len, counter, acc, previous) do
    block = :crypto.mac(:hmac, :sha512, prk, previous <> info <> <<counter>>)
    expand_hkdf(prk, info, len, counter + 1, acc <> block, block)
  end
end