defmodule Age do
  @moduledoc """
  Elixir implementation of the age encryption format.
  """

  alias Age.X25519

  @version "v1"

  defp file_key() do
    :crypto.strong_rand_bytes(16)
  end

  defp version_line() do
    "age-encryption.org/#{@version}"
  end

  @doc """
  Generates a new identity keypair and formats it in Bech32.
  """
  @spec generate_identity() :: {String.t(), String.t()}
  def generate_identity() do
    {pub, priv} = X25519.generate_keypair()

    {X25519.encode_public_key(pub), X25519.encode_private_key(priv)}
  end

  @doc """
  Encrypts the given data for the given list of recipients.
  """
  @spec encrypt(binary(), [String.t()]) :: {:ok, binary()} | {:error, String.t()}
  def encrypt(data, recipients) do
    fk = file_key()

    header_recp = recipients
    |> Enum.map(fn recipient ->
      X25519.decode_public_key(recipient)
    end)
    |> Enum.map(fn recipient ->
      X25519.stanza(recipient, fk)
    end)
    |> Enum.map(fn {header, body} ->
      "-> " <> header <> "\n" <> body
    end)
    |> Enum.join("\n")

    header = version_line() <> "\n" <> header_recp <> "\n" <> "---"

    header_signature = derive_header_signature(header, fk)

    signed_header = header <> " " <> header_signature

    data_chunks = chunk(64 * 1024, data)

    nonce = :crypto.strong_rand_bytes(16)
    payload_key = HKDF.derive(:sha256, fk, 32, nonce, "payload")

    body = data_chunks
    |> Enum.with_index()
    |> Enum.map(fn {chunk, index} ->
      is_final? = index == length(data_chunks) - 1
      nonce = get_chunk_nonce(index, is_final?)
      encrypt_chunk(chunk, payload_key, nonce)
    end)
    |> Enum.join(<<>>)

    {:ok, signed_header <> "\n" <> nonce <> body}
  end

  @doc """
  Decrypts the given data using the given identity.
  """
  @spec decrypt(binary(), String.t()) :: {:ok, binary()} | {:error, String.t()}
  def decrypt(data, identity) when is_binary(data) and is_binary(identity) do
    privkey = X25519.decode_private_key(identity)

    [header, sig_and_nonce_and_body] = :binary.split(data, "--- ")

    [header_signature, nonce_and_body] = :binary.split(sig_and_nonce_and_body, "\n")

    header_signature = String.trim(header_signature)

    [version, header_recipients] = :binary.split(header, "\n")

    ^version = version_line()

    recipients = String.split(header_recipients, "-> ", trim: true)
    |> Enum.map(fn header_recipient ->
      "X25519 " <> recipient = String.trim(header_recipient)

      [eph_share, body] = String.split(recipient, "\n", trim: true)
      |> Enum.map(&Base.decode64!(&1, padding: false))

      32 = byte_size(body)
      32 = byte_size(eph_share)

      {eph_share, body}
    end)

    fk = recipients
    |> Enum.map(fn {eph_share, body} ->
      shared_secret = :crypto.compute_key(:ecdh, eph_share, privkey, :x25519)

      pubkey = X25519.derive_public_key(privkey)

      wrap_key = X25519.derive_wrap_key(shared_secret, eph_share, pubkey)

      :libsodium_crypto_aead_chacha20poly1305.ietf_decrypt(body, <<0::96>>, wrap_key)
    end)
    |> Enum.find(&is_binary(&1))

    if fk == nil do
      {:error, "no matching recipient"}
    else
      # verify header signature
      input = header <> "---"
      expected_header_signature = derive_header_signature(input, fk)
      ^expected_header_signature = header_signature

      <<nonce :: bytes-16, body :: binary>> = nonce_and_body

      data_chunks = chunk(64 * 1024, body)

      payload_key = HKDF.derive(:sha256, fk, 32, nonce, "payload")

      plaintext = data_chunks
      |> Enum.with_index()
      |> Enum.map(fn {chunk, index} ->
        is_final? = index == length(data_chunks) - 1
        nonce = get_chunk_nonce(index, is_final?)
        decrypt_chunk(chunk, payload_key, nonce)
      end)
      |> Enum.join(<<>>)

      {:ok, plaintext}
    end
  end

  defp derive_header_signature(header, fk) do
    hmac_key = HKDF.derive(:sha256, fk, 32, "", "header")

    :libsodium_crypto_auth_hmacsha256.crypto_auth_hmacsha256(header, hmac_key)
    |> Base.encode64(padding: false)
  end

  defp encrypt_chunk(data, key, nonce) do
    :libsodium_crypto_aead_chacha20poly1305.ietf_encrypt(data, nonce, key)
  end

  defp decrypt_chunk(data, key, nonce) do
    :libsodium_crypto_aead_chacha20poly1305.ietf_decrypt(data, nonce, key)
  end

  defp chunk(_size, <<>>), do: [""]
  defp chunk(size, data), do: do_chunk(size, data)

  # https://stackoverflow.com/a/5881188
  # CC BY-SA 3.0
  defp do_chunk(size, data) when byte_size(data) > size do
    {head, rest} = :erlang.split_binary(data, size)
    [head | do_chunk(size, rest)]
  end

  defp do_chunk(_size, <<>>), do: []

  defp do_chunk(_size, data), do: [data]

  defp get_chunk_nonce(index, is_final?) do
    final = if is_final?, do: 1, else: 0
    <<index :: 11*8-big, final :: 8>>
  end
end