aerospike_driver

defmodule Aerospike.Protocol.AsmMsg.Value do
  @moduledoc false

  alias Aerospike.Error
  alias Aerospike.Geo
  alias Aerospike.Protocol.AsmMsg.Operation
  alias Aerospike.Protocol.MessagePack

  @particle_null 0
  @particle_integer 1
  @particle_float 2
  @particle_string 3
  @particle_blob 4
  @particle_bool 17
  @particle_map 19
  @particle_list 20
  @particle_geojson 23
  @op_write Operation.op_write()
  @op_add Operation.op_add()
  @op_append Operation.op_append()
  @op_prepend Operation.op_prepend()

  @doc """
  Encodes an Elixir value into `{particle_type, data}` for a simple write op.

  Supported values are the current-driver write subset:
  `nil`, integers, floats, strings, booleans, lists, maps, explicit blobs via
  `{:blob, binary}`, typed geo values, and explicit GeoJSON via
  `{:geojson, binary}`.
  """
  @spec encode_value(term()) ::
          {:ok, {non_neg_integer(), binary()}} | {:error, Error.t()}
  def encode_value(nil), do: {:ok, {@particle_null, <<>>}}

  def encode_value(value) when is_integer(value) do
    {:ok, {@particle_integer, <<value::64-signed-big>>}}
  end

  def encode_value(value) when is_float(value) do
    {:ok, {@particle_float, <<value::64-float-big>>}}
  end

  def encode_value(value) when is_binary(value) do
    {:ok, {@particle_string, value}}
  end

  def encode_value(true), do: {:ok, {@particle_bool, <<1>>}}
  def encode_value(false), do: {:ok, {@particle_bool, <<0>>}}

  def encode_value({:blob, value}) when is_binary(value) do
    {:ok, {@particle_blob, value}}
  end

  def encode_value(%Geo.Point{} = point), do: encode_value({:geojson, Geo.to_json(point)})
  def encode_value(%Geo.Polygon{} = polygon), do: encode_value({:geojson, Geo.to_json(polygon)})
  def encode_value(%Geo.Circle{} = circle), do: encode_value({:geojson, Geo.to_json(circle)})

  def encode_value({:geojson, json}) when is_binary(json) do
    {:ok, {@particle_geojson, <<0::8, 0::16-big, json::binary>>}}
  end

  def encode_value(%_{} = value), do: unsupported_write_particle(value)

  def encode_value(value) when is_list(value) do
    with {:ok, packed} <- pack_collection_value(value) do
      {:ok, {@particle_list, MessagePack.pack!(packed)}}
    end
  end

  def encode_value(%{} = value) do
    with {:ok, packed} <- pack_collection_value(value) do
      {:ok, {@particle_map, MessagePack.pack!(packed)}}
    end
  end

  def encode_value(value) do
    unsupported_write_particle(value)
  end

  @doc """
  Builds write operations from a bin map.

  Keys may be atoms or strings; atom keys are converted with `Atom.to_string/1`.
  """
  @spec encode_bin_operations(%{optional(atom() | String.t()) => term()}) :: [Operation.t()]
  def encode_bin_operations(bins) when is_map(bins) do
    encode_bin_operations(bins, Operation.op_write())
  end

  @doc """
  Builds operations from a bin map with the specified operation type.
  """
  @spec encode_bin_operations(%{optional(atom() | String.t()) => term()}, non_neg_integer()) ::
          [Operation.t()]
  def encode_bin_operations(bins, op_type) when is_map(bins) and is_integer(op_type) do
    bins
    |> Enum.map(fn
      {k, v} when is_atom(k) -> {Atom.to_string(k), v}
      {k, v} when is_binary(k) -> {k, v}
      {k, _v} -> raise ArgumentError, "bin name must be a string or atom, got: #{inspect(k)}"
    end)
    |> Enum.sort_by(fn {name, _} -> name end)
    |> Enum.map(fn {name, value} ->
      {:ok, {particle_type, data}} = encode_value_for_op(op_type, name, value)

      %Operation{op_type: op_type, particle_type: particle_type, bin_name: name, data: data}
    end)
  end

  defp encode_value_for_op(op_type, _name, value) when op_type == @op_write do
    encode_value(value)
  end

  defp encode_value_for_op(op_type, name, value) when op_type == @op_add do
    unless is_integer(value) do
      raise ArgumentError, "add requires integer values, got #{inspect(value)} for bin #{name}"
    end

    encode_value(value)
  end

  defp encode_value_for_op(op_type, name, value)
       when op_type in [@op_append, @op_prepend] do
    unless is_binary(value) do
      op_name = if op_type == @op_append, do: "append", else: "prepend"

      raise ArgumentError,
            "#{op_name} requires string values, got #{inspect(value)} for bin #{name}"
    end

    encode_value(value)
  end

  defp pack_collection_value(value) when is_list(value) do
    value
    |> Enum.reduce_while({:ok, []}, fn item, {:ok, acc} ->
      case pack_collection_value(item) do
        {:ok, packed} -> {:cont, {:ok, [packed | acc]}}
        {:error, %Error{}} = error -> {:halt, error}
      end
    end)
    |> case do
      {:ok, packed} -> {:ok, Enum.reverse(packed)}
      {:error, %Error{}} = error -> error
    end
  end

  defp pack_collection_value(%{} = value) do
    value
    |> Enum.reduce_while({:ok, []}, fn {key, item}, {:ok, acc} ->
      with {:ok, packed_key} <- pack_collection_map_key(key),
           {:ok, packed_value} <- pack_collection_value(item) do
        {:cont, {:ok, [{packed_key, packed_value} | acc]}}
      else
        {:error, %Error{}} = error -> {:halt, error}
      end
    end)
    |> case do
      {:ok, packed} -> {:ok, Map.new(packed)}
      {:error, %Error{}} = error -> error
    end
  end

  defp pack_collection_value(nil), do: {:ok, nil}
  defp pack_collection_value(true), do: {:ok, true}
  defp pack_collection_value(false), do: {:ok, false}
  defp pack_collection_value(value) when is_integer(value), do: {:ok, value}
  defp pack_collection_value(value) when is_float(value), do: {:ok, value}
  defp pack_collection_value(value) when is_binary(value), do: {:ok, {:particle_string, value}}

  defp pack_collection_value(value), do: unsupported_write_particle(value)

  defp pack_collection_map_key(key) when is_binary(key), do: {:ok, {:particle_string, key}}

  defp pack_collection_map_key(key) when is_atom(key),
    do: {:ok, {:particle_string, Atom.to_string(key)}}

  defp pack_collection_map_key(key) do
    {:error,
     Error.from_result_code(:invalid_argument,
       message: "map key must be a string or atom, got: #{inspect(key)}"
     )}
  end

  defp unsupported_write_particle(value) do
    {:error,
     Error.from_result_code(:invalid_argument,
       message:
         "unsupported write particle #{inspect(value)}; supported values: nil, integer, float, binary, boolean, list, map, {:blob, binary}, Aerospike.Geo structs, {:geojson, binary}"
     )}
  end

  @doc """
  Decodes wire particle bytes to an Elixir term.

  Returns `{:ok, term}` on success or `{:error, Error.t()}` on failure.
  Unsupported particle types return `{:ok, {:raw, particle_type, data}}`.
  """
  @spec decode_value(non_neg_integer(), binary()) :: {:ok, term()} | {:error, Error.t()}
  def decode_value(@particle_null, <<>>), do: {:ok, nil}

  def decode_value(@particle_integer, <<n::64-signed-big>>), do: {:ok, n}

  def decode_value(@particle_float, <<f::64-float-big>>), do: {:ok, f}

  def decode_value(@particle_string, data) when is_binary(data), do: {:ok, data}

  def decode_value(@particle_blob, data) when is_binary(data), do: {:ok, {:blob, data}}

  def decode_value(@particle_bool, <<0>>), do: {:ok, false}
  def decode_value(@particle_bool, <<1>>), do: {:ok, true}

  def decode_value(@particle_map, data) when is_binary(data) do
    decode_message_pack_particle(data, "map")
  end

  def decode_value(@particle_list, data) when is_binary(data) do
    decode_message_pack_particle(data, "list")
  end

  def decode_value(@particle_geojson, <<0::8, ncells::16-big, rest::binary>> = data)
      when byte_size(rest) >= ncells * 8 do
    cell_bytes = ncells * 8
    <<_cells::binary-size(cell_bytes), json::binary>> = rest

    if geojson_start?(json) do
      {:ok, Geo.from_json(json)}
    else
      {:ok, Geo.from_json(data)}
    end
  end

  def decode_value(@particle_geojson, data) when is_binary(data) do
    {:ok, Geo.from_json(data)}
  end

  def decode_value(particle_type, data) when is_integer(particle_type) and is_binary(data) do
    {:ok, {:raw, particle_type, data}}
  end

  defp geojson_start?(<<"{", _rest::binary>>), do: true
  defp geojson_start?(<<"[", _rest::binary>>), do: true
  defp geojson_start?(_data), do: false

  defp decode_message_pack_particle(data, type) do
    case MessagePack.unpack(data) do
      {:ok, {value, <<>>}} ->
        {:ok, value}

      {:ok, {_value, rest}} ->
        {:error,
         Error.from_result_code(:parse_error,
           message: "MessagePack #{type}: trailing bytes (#{byte_size(rest)})"
         )}

      {:error, :invalid_msgpack} ->
        {:error, Error.from_result_code(:parse_error, message: "invalid MessagePack #{type}")}
    end
  end
end