mongodb_driver

defmodule Mongo.Cursor do
  @moduledoc false

  import Record, only: [defrecordp: 2]

  @type t :: %__MODULE__{
    conn: Mongo.conn,
    coll: Mongo.collection,
    query: BSON.document,
    opts: Keyword.t
  }

  defstruct [:conn, :coll, :query, :opts]

  defimpl Enumerable do

    defrecordp :state, [:conn, :cursor, :coll, :buffer]

    def reduce(%{conn: conn, coll: coll, query: query, opts: opts}, acc, reduce_fun) do
      start_fun = start_fun(conn, coll, query, opts)
      next_fun  = next_fun(opts)
      after_fun = after_fun(opts)

      Stream.resource(start_fun, next_fun, after_fun).(acc, reduce_fun)
    end

    defp start_fun(conn, coll, query, opts) do
      opts = Keyword.put(opts, :batch_size, -1)

      fn ->
        case Mongo.direct_command(conn, query, opts) do
          {:ok, %{"ok" => ok,
                  "cursor" => %{
                    "id" => cursor,
                    "ns" => coll,
                    "firstBatch" => docs}}} when ok == 1 -> state(conn: conn, cursor: cursor, coll: coll, buffer: docs)
          {:error, error}                                -> raise error
        end
      end
    end

    defp next_fun(opts) do
      fn
        state(buffer: [], cursor: 0) = state ->  {:halt, state}

        state(buffer: [], conn: conn, cursor: cursor, coll: coll) = state ->
          case get_more(conn, only_coll(coll), cursor, opts) do
            {:ok, %{cursor_id: cursor, docs: []}}   -> {:halt, state(state, cursor: cursor)}
            {:ok, %{cursor_id: cursor, docs: docs}} -> {docs, state(state, cursor: cursor)}
            {:error, error}                         -> raise error
          end

        state(buffer: buffer) = state -> {buffer, state(state, buffer: [])}
      end
    end

    @doc """
      Calls the GetCore-Command
      See https://github.com/mongodb/specifications/blob/master/source/find_getmore_killcursors_commands.rst
    """
    def get_more(conn, coll, cursor, opts) do

      query = [
        {"getMore", cursor},
        {"collection", coll},
        {"batchSize", 2}, ##opts[:batch_size]},
        {"maxTimeMS", opts[:max_time]}
      ]

      query = filter_nils(query)

      with {:ok, %{"cursor" => %{ "id" => cursor_id, "nextBatch" => docs}, "ok" => ok}} when ok == 1 <- Mongo.direct_command(conn, query, opts) do
        {:ok, %{cursor_id: cursor_id, docs: docs}}
      end

    end

    @doc """
      Calls the KillCursors-Command
      See https://github.com/mongodb/specifications/blob/master/source/find_getmore_killcursors_commands.rst
    """
    def kill_cursors(conn, coll, cursor_ids, opts) do

      query = [
        {"killCursors", coll},
        {"cursors", cursor_ids}
      ]

      query = filter_nils(query)

      with {:ok, %{"cursorsAlive" => [],
                   "cursorsNotFound" => [],
                   "cursorsUnknown" => [],
                   "ok" => ok}} when ok == 1 <- Mongo.direct_command(conn, query, opts) do
        :ok
      end
    end

    defp filter_nils(keyword) when is_list(keyword) do
      Enum.reject(keyword, fn {_key, value} -> is_nil(value) end)
    end

    defp filter_nils(map) when is_map(map) do
      Enum.reject(map, fn {_key, value} -> is_nil(value) end)
      |> Enum.into(%{})
    end

    defp after_fun(opts) do
      fn
        state(cursor: 0)                              -> :ok
        state(cursor: cursor, coll: coll, conn: conn) -> kill_cursors(conn, only_coll(coll), [cursor], opts)
      end
    end

    defp only_coll(coll) do
      [_db, coll] = String.split(coll, ".", parts: 2)
      coll
    end

    # we cannot determinstically slice, so tell Enumerable to
    # fall back on brute force
    def slice(_cursor), do: { :error, __MODULE__ }
    def count(_stream), do: {:error, __MODULE__}
    def member?(_stream, _term), do: {:error, __MODULE__}
  end
end