defmodule AirPlay.V2.Plist do
  @moduledoc """
  Minimal binary plist writer for AirPlay 2 SETUP bodies.

  Supports the primitive shapes AP2 setup needs: maps with string keys, lists,
  binaries, strings, booleans, floats and non-negative integers.
  """

  import Bitwise, only: [&&&: 2, <<<: 2, |||: 2, >>>: 2]

  @spec encode!(term()) :: binary()
  def encode!(value) do
    {root_ref, objects} = flatten(value, [])
    object_count = length(objects)
    ref_size = int_size(max(object_count - 1, 0))

    object_table =
      objects
      |> Enum.reverse()
      |> Enum.map(&encode_object(&1, ref_size))

    object_table = IO.iodata_to_binary(object_table)
    offsets = offsets_from_lengths(Enum.reverse(objects), 8, ref_size)
    offset_table_offset = 8 + byte_size(object_table)
    offset_size = int_size(offset_table_offset)

    offset_table =
      offsets
      |> Enum.map(&uint(&1, offset_size))
      |> IO.iodata_to_binary()

    trailer =
      <<0::48, offset_size, ref_size, object_count::64, root_ref::64, offset_table_offset::64>>

    "bplist00" <> object_table <> offset_table <> trailer
  end

  @doc "Encode a term as a binary plist."
  @spec encode(term()) :: binary()
  def encode(value), do: encode!(value)

  @doc "Decode the subset of binary plist values emitted by AirPlay receivers."
  @spec decode!(binary()) :: term()
  def decode!(<<"bplist00", _rest::binary>> = data) do
    size = byte_size(data)
    trailer = binary_part(data, size - 32, 32)

    <<0::48, offset_size, ref_size, object_count::64, root_ref::64, offset_table_offset::64>> =
      trailer

    offsets =
      for index <- 0..(object_count - 1) do
        offset_table_offset
        |> Kernel.+(index * offset_size)
        |> then(&binary_part(data, &1, offset_size))
        |> read_uint()
      end

    {value, _cache} = decode_ref(data, offsets, ref_size, root_ref, %{})
    value
  end

  def decode!(data) when is_binary(data), do: raise(ArgumentError, "not a binary plist")

  @doc "Decode a binary plist into `{:ok, term}` or `{:error, reason}`."
  @spec decode(binary()) :: {:ok, term()} | {:error, term()}
  def decode(<<"bplist00", _rest::binary>> = data), do: {:ok, decode!(data)}
  def decode(data) when is_binary(data), do: {:error, :not_bplist}

  defp decode_ref(_data, _offsets, _ref_size, ref, cache) when is_map_key(cache, ref) do
    {Map.fetch!(cache, ref), cache}
  end

  defp decode_ref(data, offsets, ref_size, ref, cache) do
    offset = Enum.at(offsets, ref)
    marker = :binary.at(data, offset)
    type = marker >>> 4
    info = marker &&& 0x0F

    {value, cache} = decode_typed_value(type, info, data, offset, offsets, ref_size, cache)

    {value, Map.put(cache, ref, value)}
  end

  defp decode_typed_value(0x0, info, _data, _offset, _offsets, _ref_size, cache) do
    {simple_value(info), cache}
  end

  defp decode_typed_value(0x1, info, data, offset, _offsets, _ref_size, cache) do
    size = 1 <<< info
    {data |> binary_part(offset + 1, size) |> read_int(size), cache}
  end

  defp decode_typed_value(0x2, info, data, offset, _offsets, _ref_size, cache) do
    size = 1 <<< info

    value =
      case binary_part(data, offset + 1, size) do
        <<float::float-big-32>> -> float
        <<float::float-big-64>> -> float
      end

    {value, cache}
  end

  defp decode_typed_value(0x4, info, data, offset, _offsets, _ref_size, cache) do
    {count, value_offset} = decode_count(data, offset, info)
    {{:data, binary_part(data, value_offset, count)}, cache}
  end

  defp decode_typed_value(0x5, info, data, offset, _offsets, _ref_size, cache) do
    {count, value_offset} = decode_count(data, offset, info)
    {binary_part(data, value_offset, count), cache}
  end

  defp decode_typed_value(0x6, info, data, offset, _offsets, _ref_size, cache) do
    {count, value_offset} = decode_count(data, offset, info)

    value =
      data
      |> binary_part(value_offset, count * 2)
      |> :unicode.characters_to_binary({:utf16, :big}, :utf8)

    {value, cache}
  end

  defp decode_typed_value(0xA, info, data, offset, offsets, ref_size, cache) do
    {count, refs_offset} = decode_count(data, offset, info)
    refs = read_refs(data, refs_offset, ref_size, count)
    decode_refs(refs, data, offsets, ref_size, cache)
  end

  defp decode_typed_value(0xD, info, data, offset, offsets, ref_size, cache) do
    {count, refs_offset} = decode_count(data, offset, info)
    key_refs = read_refs(data, refs_offset, ref_size, count)
    value_refs = read_refs(data, refs_offset + count * ref_size, ref_size, count)
    {keys, cache} = decode_refs(key_refs, data, offsets, ref_size, cache)
    {values, cache} = decode_refs(value_refs, data, offsets, ref_size, cache)
    {keys |> Enum.zip(values) |> Map.new(), cache}
  end

  defp decode_refs(refs, data, offsets, ref_size, cache) do
    Enum.map_reduce(refs, cache, fn child_ref, cache ->
      decode_ref(data, offsets, ref_size, child_ref, cache)
    end)
  end

  defp simple_value(0x0), do: nil
  defp simple_value(0x8), do: false
  defp simple_value(0x9), do: true

  defp decode_count(_data, offset, info) when info < 0xF, do: {info, offset + 1}

  defp decode_count(data, offset, 0xF) do
    int_offset = offset + 1
    marker = :binary.at(data, int_offset)
    0x1 = marker >>> 4
    size = 1 <<< (marker &&& 0x0F)
    count = data |> binary_part(int_offset + 1, size) |> read_uint()
    {count, int_offset + 1 + size}
  end

  defp read_refs(_data, _offset, _ref_size, 0), do: []

  defp read_refs(data, offset, ref_size, count) do
    for index <- 0..(count - 1) do
      data
      |> binary_part(offset + index * ref_size, ref_size)
      |> read_uint()
    end
  end

  defp flatten(value, objects) when is_map(value) do
    pairs = value |> Enum.sort_by(fn {key, _value} -> to_string(key) end)

    {key_refs, objects} =
      Enum.map_reduce(pairs, objects, fn {key, _value}, acc -> flatten(to_string(key), acc) end)

    {value_refs, objects} =
      Enum.map_reduce(pairs, objects, fn {_key, value}, acc -> flatten(value, acc) end)

    add_object({:dict, key_refs, value_refs}, objects)
  end

  defp flatten(value, objects) when is_list(value) do
    {refs, objects} = Enum.map_reduce(value, objects, &flatten/2)
    add_object({:array, refs}, objects)
  end

  defp flatten({:data, value}, objects) when is_binary(value),
    do: add_object({:data, value}, objects)

  defp flatten(value, objects) when is_binary(value) do
    if String.valid?(value),
      do: add_object({:string, value}, objects),
      else: add_object({:data, value}, objects)
  end

  defp flatten(value, objects) when is_boolean(value), do: add_object({:bool, value}, objects)
  defp flatten(value, objects) when is_float(value), do: add_object({:real, value}, objects)

  defp flatten(value, objects) when is_integer(value),
    do: add_object({:int, value}, objects)

  defp add_object(object, objects), do: {length(objects), [object | objects]}

  defp encode_object({:bool, false}, _ref_size), do: <<0x08>>
  defp encode_object({:bool, true}, _ref_size), do: <<0x09>>
  defp encode_object({:int, value}, _ref_size), do: encode_int(value)
  defp encode_object({:real, value}, _ref_size), do: encode_real(value)
  defp encode_object({:data, value}, _ref_size), do: [marker(0x40, byte_size(value)), value]
  defp encode_object({:string, value}, _ref_size), do: [marker(0x50, byte_size(value)), value]

  defp encode_object({:array, refs}, ref_size) do
    [marker(0xA0, length(refs)), Enum.map(refs, &uint(&1, ref_size))]
  end

  defp encode_object({:dict, key_refs, value_refs}, ref_size) do
    [
      marker(0xD0, length(key_refs)),
      Enum.map(key_refs, &uint(&1, ref_size)),
      Enum.map(value_refs, &uint(&1, ref_size))
    ]
  end

  defp encode_int(value) do
    size = int_size(value)
    exponent = round(:math.log2(size))
    [<<0x10 ||| exponent>>, int(value, size)]
  end

  defp encode_real(value), do: <<0x23, value::float-big-size(64)>>

  defp marker(base, length) when length < 15, do: <<base ||| length>>
  defp marker(base, length), do: [<<base ||| 0x0F>>, encode_int(length)]

  defp offsets_from_lengths(objects, start, ref_size) do
    {_offset, offsets} =
      objects
      |> Enum.map(&encode_object(&1, ref_size))
      |> Enum.map(&IO.iodata_to_binary/1)
      |> Enum.reduce({start, []}, fn encoded, {offset, acc} ->
        {offset + byte_size(encoded), [offset | acc]}
      end)

    Enum.reverse(offsets)
  end

  defp int_size(value) when value < 0, do: 8
  defp int_size(value) when value <= 0xFF, do: 1
  defp int_size(value) when value <= 0xFFFF, do: 2
  defp int_size(value) when value <= 0xFFFFFFFF, do: 4
  defp int_size(_value), do: 8

  defp int(value, size) when value < 0, do: <<value::signed-big-integer-size(size * 8)>>
  defp int(value, size), do: uint(value, size)

  defp uint(value, 1), do: <<value::8>>
  defp uint(value, 2), do: <<value::16>>
  defp uint(value, 4), do: <<value::32>>
  defp uint(value, 8), do: <<value::64>>

  defp read_int(binary, 8), do: binary |> read_signed_64()
  defp read_int(binary, _size), do: read_uint(binary)

  defp read_signed_64(<<value::signed-big-64>>), do: value

  defp read_uint(binary), do: :binary.decode_unsigned(binary)
end