defmodule AirPlay.V2.TLV do
  @moduledoc """
  AirPlay 2 TLV8 encoder/decoder.

  Values longer than 255 bytes are split into repeated type entries. Decoding
  rejoins repeated chunks for the same type, matching HomeKit/AirPlay TLV8 use.
  """

  @types %{
    method: 0,
    identifier: 1,
    salt: 2,
    public_key: 3,
    proof: 4,
    encrypted_data_with_tag: 5,
    state: 6,
    error: 7,
    signature: 10,
    flags: 19
  }

  @names Map.new(@types, fn {name, id} -> {id, name} end)

  @type type :: atom() | 0..255
  @type pair :: {type(), binary() | iodata() | non_neg_integer()}

  @doc "Encode a TLV map or key/value list."
  @spec encode(map() | [pair()]) :: binary()
  def encode(values) when is_map(values), do: values |> Map.to_list() |> encode()

  def encode(values) when is_list(values) do
    values
    |> Enum.map(fn {type, value} -> encode_value(type_id(type), normalize_value(value)) end)
    |> IO.iodata_to_binary()
  end

  @doc "Decode TLV8 into a map keyed by known atoms or raw integer types."
  @spec decode(binary()) :: map() | {:error, term()}
  def decode(binary) when is_binary(binary) do
    case decode_result(binary) do
      {:ok, map} -> map
      {:error, _reason} = error -> error
    end
  end

  @doc "Decode TLV8 into an explicit `{:ok, map}` / `{:error, reason}` tuple."
  @spec decode_result(binary()) :: {:ok, map()} | {:error, term()}
  def decode_result(binary) when is_binary(binary), do: do_decode(binary, %{})

  @doc "Return the numeric TLV type id for an atom or integer type."
  @spec type_id(type()) :: 0..255
  def type_id(type) when is_atom(type), do: Map.fetch!(@types, type)
  def type_id(type) when is_integer(type) and type in 0..255, do: type

  @doc "Return the canonical atom name for a known type id, else the raw id."
  @spec type_name(0..255) :: atom() | 0..255
  def type_name(type) when is_integer(type), do: Map.get(@names, type, type)

  defp encode_value(type, value) when byte_size(value) <= 255 do
    <<type, byte_size(value), value::binary>>
  end

  defp encode_value(type, value) do
    <<chunk::binary-size(255), rest::binary>> = value
    [<<type, 255, chunk::binary>>, encode_value(type, rest)]
  end

  defp do_decode(<<>>, acc), do: {:ok, acc}

  defp do_decode(<<type, len, rest::binary>>, acc) when byte_size(rest) >= len do
    <<value::binary-size(^len), tail::binary>> = rest
    key = type_name(type)
    acc = Map.update(acc, key, value, &(&1 <> value))
    do_decode(tail, acc)
  end

  defp do_decode(_binary, _acc), do: {:error, :truncated_tlv}

  defp normalize_value(value) when is_binary(value), do: value
  defp normalize_value(value) when is_integer(value) and value in 0..255, do: <<value>>
  defp normalize_value(value), do: IO.iodata_to_binary(value)
end