aerospike_driver

defmodule Aerospike.Command.ScanOps do
  @moduledoc false

  alias Aerospike.Command.NodePartitions
  alias Aerospike.Command.PartitionTracker
  alias Aerospike.Command.ScanOps.PageRunner
  alias Aerospike.Command.ScanOps.StreamRunner
  alias Aerospike.Cursor
  alias Aerospike.Error
  alias Aerospike.ExecuteTask
  alias Aerospike.Page
  alias Aerospike.PartitionFilter
  alias Aerospike.Protocol.AsmMsg.Operation
  alias Aerospike.Query
  alias Aerospike.Query.AggregateReducer
  alias Aerospike.Scan

  @spec stream(GenServer.server(), Scan.t() | Query.t(), keyword()) ::
          {:ok, Enumerable.t()} | {:error, Error.t()}
  def stream(tender, scannable, opts \\ []) when is_list(opts) do
    with {:ok, node_name, opts2} <- normalize_node_opt(opts) do
      {cursor, opts3} = Keyword.pop(opts2, :cursor)

      with {:ok, scannable2} <- apply_optional_cursor(scannable, cursor) do
        stream_with_node(tender, scannable2, node_name, opts3)
      end
    end
  end

  @spec stream_node(GenServer.server(), String.t(), Scan.t() | Query.t(), keyword()) ::
          {:ok, Enumerable.t()} | {:error, Error.t()}
  def stream_node(tender, node_name, scannable, opts \\ [])
      when is_binary(node_name) and is_list(opts) do
    with {:ok, opts2} <- put_target_node_opt(opts, node_name) do
      stream(tender, scannable, opts2)
    end
  end

  @spec all(GenServer.server(), Scan.t(), keyword()) ::
          {:ok, [Aerospike.Record.t()]} | {:error, Error.t()}
  def all(tender, %Scan{} = scan, opts \\ []) when is_list(opts) do
    collect(tender, scan, opts, &Enum.to_list/1)
  end

  @spec all_node(GenServer.server(), String.t(), Scan.t(), keyword()) ::
          {:ok, [Aerospike.Record.t()]} | {:error, Error.t()}
  def all_node(tender, node_name, %Scan{} = scan, opts \\ [])
      when is_binary(node_name) and is_list(opts) do
    with {:ok, opts2} <- put_target_node_opt(opts, node_name) do
      all(tender, scan, opts2)
    end
  end

  @spec scan_page(GenServer.server(), Scan.t(), keyword()) ::
          {:ok, Page.t()} | {:error, Error.t()}
  def scan_page(tender, %Scan{} = scan, opts \\ []) when is_list(opts) do
    with {:ok, node_name, opts2} <- normalize_node_opt(opts) do
      {cursor, opts3} = Keyword.pop(opts2, :cursor)

      with {:ok, scan2} <- apply_optional_cursor(scan, cursor),
           :ok <- require_max_records(scan2) do
        scan_page_with_node(tender, scan2, node_name, opts3)
      end
    end
  end

  @spec scan_page_node(GenServer.server(), String.t(), Scan.t(), keyword()) ::
          {:ok, Page.t()} | {:error, Error.t()}
  def scan_page_node(tender, node_name, %Scan{} = scan, opts \\ [])
      when is_binary(node_name) and is_list(opts) do
    with {:ok, opts2} <- put_target_node_opt(opts, node_name) do
      scan_page(tender, scan, opts2)
    end
  end

  @spec count(GenServer.server(), Scan.t() | Query.t(), keyword()) ::
          {:ok, non_neg_integer()} | {:error, Error.t()}
  def count(tender, scannable, opts \\ []) when is_list(opts) do
    collect(tender, scannable, opts, &Enum.count/1)
  end

  @spec count_node(GenServer.server(), String.t(), Scan.t() | Query.t(), keyword()) ::
          {:ok, non_neg_integer()} | {:error, Error.t()}
  def count_node(tender, node_name, scannable, opts \\ [])
      when is_binary(node_name) and is_list(opts) do
    with {:ok, opts2} <- put_target_node_opt(opts, node_name) do
      count(tender, scannable, opts2)
    end
  end

  @spec query_stream(GenServer.server(), Query.t(), keyword()) ::
          {:ok, Enumerable.t()} | {:error, Error.t()}
  def query_stream(tender, %Query{} = query, opts \\ []) when is_list(opts) do
    stream(tender, query, opts)
  end

  @spec query_stream_node(GenServer.server(), String.t(), Query.t(), keyword()) ::
          {:ok, Enumerable.t()} | {:error, Error.t()}
  def query_stream_node(tender, node_name, %Query{} = query, opts \\ [])
      when is_binary(node_name) and is_list(opts) do
    stream_node(tender, node_name, query, opts)
  end

  @spec query_aggregate(GenServer.server(), Query.t(), String.t(), String.t(), list(), keyword()) ::
          {:ok, Enumerable.t()} | {:error, Error.t()}
  def query_aggregate(tender, %Query{} = query, package, function, args, opts \\ [])
      when is_binary(package) and is_binary(function) and is_list(args) and is_list(opts) do
    StreamRunner.query_aggregate(tender, query, opts, package, function, args)
  end

  @spec query_aggregate_result(
          GenServer.server(),
          Query.t(),
          String.t(),
          String.t(),
          list(),
          keyword()
        ) :: {:ok, term() | nil} | {:error, Error.t()}
  def query_aggregate_result(tender, %Query{} = query, package, function, args, opts \\ [])
      when is_binary(package) and is_binary(function) and is_list(args) and is_list(opts) do
    with {:ok, reducer} <- AggregateReducer.prepare(package, function, args, opts),
         {:ok, stream} <-
           query_aggregate(tender, query, package, function, args, aggregate_transport_opts(opts)) do
      AggregateReducer.run(reducer, stream)
    end
  end

  @spec query_execute(GenServer.server(), Query.t(), list(), keyword()) ::
          {:ok, ExecuteTask.t()} | {:error, Error.t()}
  def query_execute(tender, %Query{} = query, ops, opts \\ [])
      when is_list(ops) and is_list(opts) do
    with {:ok, built_ops} <- build_background_operations(ops),
         {:ok, node_name, opts2} <- normalize_node_opt(opts) do
      query_execute_with_node(tender, query, built_ops, node_name, opts2)
    end
  end

  @spec query_execute_node(GenServer.server(), String.t(), Query.t(), list(), keyword()) ::
          {:ok, ExecuteTask.t()} | {:error, Error.t()}
  def query_execute_node(tender, node_name, %Query{} = query, ops, opts \\ [])
      when is_binary(node_name) and is_list(ops) and is_list(opts) do
    with {:ok, opts2} <- put_target_node_opt(opts, node_name) do
      query_execute(tender, query, ops, opts2)
    end
  end

  @spec query_udf(GenServer.server(), Query.t(), String.t(), String.t(), list(), keyword()) ::
          {:ok, ExecuteTask.t()} | {:error, Error.t()}
  def query_udf(tender, %Query{} = query, package, function, args, opts \\ [])
      when is_binary(package) and is_binary(function) and is_list(args) and is_list(opts) do
    with {:ok, node_name, opts2} <- normalize_node_opt(opts) do
      query_udf_with_node(tender, query, package, function, args, node_name, opts2)
    end
  end

  @spec query_udf_node(
          GenServer.server(),
          String.t(),
          Query.t(),
          String.t(),
          String.t(),
          list(),
          keyword()
        ) :: {:ok, ExecuteTask.t()} | {:error, Error.t()}
  def query_udf_node(tender, node_name, %Query{} = query, package, function, args, opts \\ [])
      when is_binary(node_name) and is_binary(package) and is_binary(function) and is_list(args) and
             is_list(opts) do
    with {:ok, opts2} <- put_target_node_opt(opts, node_name) do
      query_udf(tender, query, package, function, args, opts2)
    end
  end

  @spec query_all(GenServer.server(), Query.t(), keyword()) ::
          {:ok, [Aerospike.Record.t()]} | {:error, Error.t()}
  def query_all(tender, %Query{} = query, opts \\ []) when is_list(opts) do
    with :ok <- require_max_records(query),
         {:ok, node_name, opts2} <- normalize_node_opt(opts) do
      query_all_with_node(tender, query, node_name, opts2)
    end
  end

  @spec query_all_node(GenServer.server(), String.t(), Query.t(), keyword()) ::
          {:ok, [Aerospike.Record.t()]} | {:error, Error.t()}
  def query_all_node(tender, node_name, %Query{} = query, opts \\ [])
      when is_binary(node_name) and is_list(opts) do
    with {:ok, opts2} <- put_target_node_opt(opts, node_name) do
      query_all(tender, query, opts2)
    end
  end

  @spec query_count(GenServer.server(), Query.t(), keyword()) ::
          {:ok, non_neg_integer()} | {:error, Error.t()}
  def query_count(tender, %Query{} = query, opts \\ []) when is_list(opts) do
    count(tender, query, opts)
  end

  @spec query_count_node(GenServer.server(), String.t(), Query.t(), keyword()) ::
          {:ok, non_neg_integer()} | {:error, Error.t()}
  def query_count_node(tender, node_name, %Query{} = query, opts \\ [])
      when is_binary(node_name) and is_list(opts) do
    count_node(tender, node_name, query, opts)
  end

  @spec query_page(GenServer.server(), Query.t(), keyword()) ::
          {:ok, Page.t()} | {:error, Error.t()}
  def query_page(tender, %Query{} = query, opts \\ []) when is_list(opts) do
    with {:ok, node_name, opts2} <- normalize_node_opt(opts) do
      {cursor, opts3} = Keyword.pop(opts2, :cursor)

      with {:ok, query2} <- apply_optional_cursor(query, cursor),
           :ok <- require_max_records(query2) do
        query_page_with_node(tender, query2, node_name, opts3)
      end
    end
  end

  @spec query_page_node(GenServer.server(), String.t(), Query.t(), keyword()) ::
          {:ok, Page.t()} | {:error, Error.t()}
  def query_page_node(tender, node_name, %Query{} = query, opts \\ [])
      when is_binary(node_name) and is_list(opts) do
    with {:ok, opts2} <- put_target_node_opt(opts, node_name) do
      query_page(tender, query, opts2)
    end
  end

  defp collect(tender, scannable, opts, fun) do
    with {:ok, node_name, opts2} <- normalize_node_opt(opts),
         {:ok, stream} <- stream_with_node(tender, scannable, node_name, opts2) do
      try do
        {:ok, fun.(stream)}
      rescue
        err in Aerospike.Error -> {:error, err}
      end
    end
  end

  defp query_all_with_node(tender, query, nil, opts), do: PageRunner.all(tender, query, opts)

  defp query_all_with_node(tender, query, node_name, opts),
    do: PageRunner.all_node(tender, node_name, query, opts)

  defp scan_page_with_node(tender, scan, nil, opts), do: PageRunner.page(tender, scan, opts)

  defp scan_page_with_node(tender, scan, node_name, opts),
    do: PageRunner.page_node(tender, node_name, scan, opts)

  defp query_page_with_node(tender, query, nil, opts), do: PageRunner.page(tender, query, opts)

  defp query_page_with_node(tender, query, node_name, opts),
    do: PageRunner.page_node(tender, node_name, query, opts)

  defp stream_with_node(tender, scannable, nil, opts),
    do: StreamRunner.stream(tender, scannable, opts)

  defp stream_with_node(tender, scannable, node_name, opts),
    do: StreamRunner.stream_node(tender, node_name, scannable, opts)

  defp query_execute_with_node(tender, query, ops, nil, opts),
    do: StreamRunner.query_execute(tender, query, ops, opts)

  defp query_execute_with_node(tender, query, ops, node_name, opts),
    do: StreamRunner.query_execute_node(tender, node_name, query, ops, opts)

  defp query_udf_with_node(tender, query, package, function, args, nil, opts),
    do: StreamRunner.query_udf(tender, query, package, function, args, opts)

  defp query_udf_with_node(tender, query, package, function, args, node_name, opts),
    do: StreamRunner.query_udf_node(tender, node_name, query, package, function, args, opts)

  defp aggregate_transport_opts(opts) do
    Keyword.drop(opts, [:source, :source_path])
  end

  defp build_background_operations(operations) do
    operations
    |> Enum.reduce_while({:ok, []}, fn operation, {:ok, acc} ->
      case build_background_operation(operation) do
        {:ok, built_operation} -> {:cont, {:ok, [built_operation | acc]}}
        {:error, %Error{}} = err -> {:halt, err}
      end
    end)
    |> case do
      {:ok, built_operations} -> {:ok, Enum.reverse(built_operations)}
      {:error, %Error{}} = err -> err
    end
  end

  defp build_background_operation(%Operation{} = operation) do
    validate_background_write_operation({:ok, operation}, operation)
  end

  defp build_background_operation({:read, bin_name} = operation) do
    {:read, normalize_bin_name(bin_name)}
    |> Operation.from_simple()
    |> validate_background_write_operation(operation)
  end

  defp build_background_operation({kind, bin_name, value} = operation)
       when kind in [:write, :add, :append, :prepend] do
    {kind, normalize_bin_name(bin_name), value}
    |> Operation.from_simple()
    |> validate_background_write_operation(operation)
  end

  defp build_background_operation(operation) do
    operation
    |> Operation.from_simple()
    |> validate_background_write_operation(operation)
  end

  defp validate_background_write_operation({:ok, %Operation{} = operation}, source) do
    if background_write_operation?(operation) do
      {:ok, operation}
    else
      invalid_background_operation(source)
    end
  end

  defp validate_background_write_operation({:error, %Error{}} = err, _source), do: err

  defp background_write_operation?(%Operation{read_header: true}), do: false

  defp background_write_operation?(%Operation{op_type: op_type}) do
    op_type in [
      Operation.op_write(),
      Operation.op_cdt_modify(),
      Operation.op_add(),
      Operation.op_exp_modify(),
      Operation.op_append(),
      Operation.op_prepend(),
      Operation.op_touch(),
      Operation.op_bit_modify(),
      Operation.op_delete(),
      Operation.op_hll_modify()
    ]
  end

  defp invalid_background_operation(operation) do
    {:error,
     Error.from_result_code(:invalid_argument,
       message:
         "background query operations must all be write operations, got: #{inspect(operation)}"
     )}
  end

  defp normalize_bin_name(bin_name) when is_atom(bin_name), do: Atom.to_string(bin_name)
  defp normalize_bin_name(bin_name), do: bin_name

  defp normalize_node_opt(opts) when is_list(opts) do
    case Keyword.pop(opts, :node) do
      {nil, opts2} ->
        {:ok, nil, opts2}

      {node_name, opts2} when is_binary(node_name) ->
        {:ok, node_name, opts2}

      {other, _opts2} ->
        {:error,
         Error.from_result_code(:parameter_error,
           message: "invalid node option: #{inspect(other)}"
         )}
    end
  end

  defp put_target_node_opt(opts, node_name) when is_list(opts) and is_binary(node_name) do
    case Keyword.has_key?(opts, :node) do
      true ->
        {:error,
         Error.from_result_code(
           :parameter_error,
           message: "node option conflicts with the positional node argument"
         )}

      false ->
        {:ok, Keyword.put(opts, :node, node_name)}
    end
  end

  defp require_max_records(%Query{max_records: n}) when is_integer(n) and n > 0, do: :ok
  defp require_max_records(%Scan{max_records: n}) when is_integer(n) and n > 0, do: :ok
  defp require_max_records(_), do: {:error, Error.from_result_code(:max_records_required)}

  defp apply_optional_cursor(scannable, nil), do: {:ok, scannable}

  defp apply_optional_cursor(scannable, %Cursor{} = cursor),
    do: {:ok, attach_cursor_partition_filter(scannable, cursor)}

  defp apply_optional_cursor(scannable, bin) when is_binary(bin) do
    case Cursor.decode(bin) do
      {:ok, cursor} -> {:ok, attach_cursor_partition_filter(scannable, cursor)}
      {:error, _} = err -> err
    end
  end

  defp apply_optional_cursor(_scannable, other) do
    {:error,
     Error.from_result_code(:parameter_error, message: "invalid cursor: #{inspect(other)}")}
  end

  defp attach_cursor_partition_filter(%Query{} = query, %Cursor{partitions: partitions}) do
    %{query | partition_filter: %{PartitionFilter.all() | partitions: partitions}}
  end

  defp attach_cursor_partition_filter(%Scan{} = scan, %Cursor{partitions: partitions}) do
    %{scan | partition_filter: %{PartitionFilter.all() | partitions: partitions}}
  end

  @doc false
  @spec allow_record_fold(
          PartitionTracker.t(),
          NodePartitions.t(),
          [term()]
        ) ::
          {PartitionTracker.t(), NodePartitions.t(), [term()]}
  def allow_record_fold(tracker, np, records) when is_list(records) do
    {tracker, np, kept_records_rev} =
      Enum.reduce(records, {tracker, np, []}, fn record, {tracker_acc, np_acc, kept_acc} ->
        case PartitionTracker.allow_record?(tracker_acc, np_acc) do
          {true, tracker2, np2} -> {tracker2, np2, [record | kept_acc]}
          {false, tracker2, np2} -> {tracker2, np2, kept_acc}
        end
      end)

    {tracker, np, Enum.reverse(kept_records_rev)}
  end
end