defmodule CljElixir.NRepl.Bencode do
  @moduledoc "Bencode encoder/decoder for nREPL wire protocol."

  # --- Encoder ---

  def encode(data) when is_binary(data) do
    "#{byte_size(data)}:#{data}"
  end

  def encode(data) when is_integer(data) do
    "i#{data}e"
  end

  def encode(data) when is_list(data) do
    "l" <> Enum.map_join(data, "", &encode/1) <> "e"
  end

  def encode(data) when is_map(data) do
    sorted = data |> Enum.sort_by(fn {k, _} -> to_string(k) end)

    "d" <>
      Enum.map_join(sorted, "", fn {k, v} ->
        encode(to_string(k)) <> encode(v)
      end) <> "e"
  end

  def encode(data) when is_atom(data) do
    encode(Atom.to_string(data))
  end

  # --- Decoder ---

  def decode(data) when is_binary(data) do
    {result, _rest} = decode_one(data)
    result
  end

  def decode_all(data) do
    decode_all(data, [])
  end

  defp decode_all("", acc), do: Enum.reverse(acc)

  defp decode_all(data, acc) do
    {result, rest} = decode_one(data)
    decode_all(rest, [result | acc])
  end

  # Public so the TCP server can call it for streaming decode
  def decode_one("i" <> rest) do
    {num_str, "e" <> rest2} = split_at_char(rest, ?e)
    {String.to_integer(num_str), rest2}
  end

  def decode_one("l" <> rest) do
    decode_list(rest, [])
  end

  def decode_one("d" <> rest) do
    decode_dict(rest, %{})
  end

  def decode_one(data) do
    # String: <length>:<data>
    {len_str, ":" <> rest} = split_at_char(data, ?:)
    len = String.to_integer(len_str)
    <<str::binary-size(len), rest2::binary>> = rest
    {str, rest2}
  end

  defp decode_list("e" <> rest, acc), do: {Enum.reverse(acc), rest}

  defp decode_list(data, acc) do
    {item, rest} = decode_one(data)
    decode_list(rest, [item | acc])
  end

  defp decode_dict("e" <> rest, acc), do: {acc, rest}

  defp decode_dict(data, acc) do
    {key, rest} = decode_one(data)
    {value, rest2} = decode_one(rest)
    decode_dict(rest2, Map.put(acc, key, value))
  end

  defp split_at_char(data, char) do
    case :binary.match(data, <<char>>) do
      {pos, 1} ->
        <<before::binary-size(pos), _::8, rest::binary>> = data
        {before, <<char, rest::binary>>}

      :nomatch ->
        {data, ""}
    end
  end
end