misha_cafex

defmodule Cafex.Consumer.OffsetManager do
  use GenServer

  require Logger

  @max_buffers 50
  @interval 500
  @storage :kafka
  @reset :latest

  defmodule State do
    @moduledoc false
    defstruct  conn: nil,  # coordinator connection
               coordinator: nil,
               partitions: nil,
               consumer_group_generation_id: nil,
               consumer_id: nil,
               consumer_group: nil,
               topic: nil,
               to_be_commit: %{},
               interval: nil,
               timer: nil,
               max_buffers: 0,
               auto_commit: true,
               reset: :latest,
               storage: nil
  end

  alias Cafex.Connection
  alias Cafex.Protocol.Offset.Request, as: OffsetRequest
  alias Cafex.Protocol.OffsetCommit
  alias Cafex.Protocol.OffsetFetch

  # ===================================================================
  # API
  # ===================================================================

  def start_link(coordinator, partitions, group, topic, opts \\ []) do
    GenServer.start_link __MODULE__, [coordinator, partitions, group, topic, opts]
  end

  def stop(pid) do
    GenServer.call pid, :stop
  end

  def update_generation_id(pid, consumer_id, generation_id) do
    GenServer.call pid, {:update_generation_id, consumer_id, generation_id}
  end

  def commit(pid, partition, offset, metadata \\ "") do
    GenServer.call pid, {:commit, partition, offset, metadata}
  end

  def fetch(pid, partition, leader_conn) do
    GenServer.call pid, {:fetch, partition, leader_conn}
  end

  def reset(pid, partition, leader_conn) do
    GenServer.call pid, {:reset, partition, leader_conn}
  end

  # ===================================================================
  # GenServer callbacks
  # ===================================================================

  def init([coordinator, partitions, group, topic, opts]) do
    state = %State{coordinator: coordinator,
                   topic: topic,
                   consumer_group: group,
                   partitions: partitions,
                   auto_commit:  Keyword.get(opts, :auto_commit, true),
                   interval:     Keyword.get(opts, :interval) || @interval,
                   max_buffers:  Keyword.get(opts, :max_buffers) || @max_buffers,
                   reset:        Keyword.get(opts, :offset_reset) || @reset,
                   storage:      Keyword.get(opts, :offset_storage) || @storage}
                 |> start_conn
    {:ok, state}
  end

  def handle_call(:stop, _from, state) do
    {:stop, :normal, :ok, state}
  end

  def handle_call({:update_generation_id, consumer_id, generation_id}, _from, state) do
    {:reply, :ok, %{state | consumer_id: consumer_id, consumer_group_generation_id: generation_id}}
  end

  def handle_call({:commit, partition, _, _}, _from, %{partitions: partitions} = state) when partition >= partitions do
    {:reply, {:error, :unknown_partition}, state}
  end
  def handle_call({:commit, partition, offset, metadata}, _from, state) do
    {reply, state} = schedule_commit(partition, offset, metadata, state)
    {:reply, reply, state}
  end

  def handle_call({:fetch, partition, _leader_conn}, _from, %{partitions: partitions} = state) when partition >= partitions do
    {:reply, {:error, :unknown_partition}, state}
  end
  def handle_call({:fetch, partition, leader_conn}, _from, %{topic: topic, reset: strategy} = state) do
    case do_fetch(partition, state) do
      {:ok, {-1, _}} ->
        {:reply, get_offset(topic, partition, leader_conn, strategy), state}
      {:ok, _} = reply ->
        {:reply, reply, state}
      {:error, :unknown_topic_or_partition} ->
        {:reply, get_offset(topic, partition, leader_conn, strategy), state}
      error ->
        {:reply, error, state}
    end
  end

  def handle_call({:reset, partition, leader_conn}, _from, %{topic: topic, reset: strategy} = state) do
    {:reply, get_offset(topic, partition, leader_conn, strategy), state}
  end

  def handle_info({:timeout, _timer, :do_commit}, %{to_be_commit: to_be_commit} = state)
  when map_size(to_be_commit) == 0 do
    {:noreply, %{state | timer: nil}}
  end
  def handle_info({:timeout, _timer, :do_commit}, state) do
    do_commit(state)
    {:noreply, %{state | to_be_commit: %{}, timer: nil}}
  end

  def terminate(_reason, state) do
    do_commit(state)
    close_conn(state)
    :ok
  end

  # ===================================================================
  # Internal functions
  # ===================================================================

  defp start_conn(%{coordinator: {host, port}} = state) do
    {:ok, pid} = Connection.start_link(host, port)
    %{state | conn: pid}
  end

  defp close_conn(%{conn: nil} = state), do: state
  defp close_conn(%{conn: pid} = state) do
    if Process.alive?(pid) do
      Connection.close(pid)
    end
    %{state | conn: nil}
  end

  defp do_fetch(partition, %{conn: conn,
                             consumer_group: group,
                             topic: topic,
                             storage: storage}) do
    request = %OffsetFetch.Request{consumer_group: group,
                                   topics: [{topic, [partition]}]}
    request = case storage do
      :zookeeper -> %{request | api_version: 0}
      :kafka     -> %{request | api_version: 1}
    end

    case Connection.request(conn, request) do
      {:ok, %{topics: [{^topic, [{^partition, offset, metadata, :no_error}]}]}} ->
        {:ok, {offset, metadata}}
      {:ok, %{topics: [{^topic, [{^partition, _, _, error}]}]}} ->
        {:error, error}
      {:error, reason} ->
        {:error, reason}
    end
  end

  defp schedule_commit(partition, offset, metadata, %{auto_commit: false} = state) do
    reply = case do_commit([{partition, offset, metadata}], state) do
      {:ok, [{^partition, :no_error}]} -> :ok
      {:ok, [{^partition, error}]} -> {:error, error}
      {:error, _reason} = error -> error
    end

    {reply, state}
  end
  defp schedule_commit(partition, offset, metadata, %{to_be_commit: to_be_commit,
                                                      max_buffers: max_buffers,
                                                      timer: timer} = state) do
    to_be_commit = Map.put(to_be_commit, partition, {offset, metadata})
    state = %{state | to_be_commit: to_be_commit}

    if timer != :sent && Map.size(to_be_commit) >= max_buffers do
      state = cancel_timer(state)
      send self, {:timeout, nil, :do_commit}
      {:ok, %{state | timer: :sent}}
    else
      {:ok, start_timer(state)}
    end
  end

  defp do_commit(%{to_be_commit: to_be_commit} = state) do
    partitions =
      to_be_commit
      |> Enum.map(fn {partition, {offset, metadata}} ->
        {partition, offset, metadata}
      end)
      |> Enum.filter(fn
        {_, offset, _} when is_integer(offset) -> true
        _ -> false
      end)
    do_commit(partitions, state)
  end

  defp do_commit([], _state) do
    {:ok, []}
  end
  defp do_commit(partitions, %{conn: conn, topic: topic, storage: storage} = state) do
    Logger.debug "Do offset commit: topic: #{inspect topic}, group: #{inspect state.consumer_group} partition offsets: #{inspect partitions},"
    request = Map.take(state, [:consumer_group,
                               :consumer_id,
                               :consumer_group_generation_id])
            |> Map.put(:topics, [{topic, partitions}])
            |> (&(struct(OffsetCommit.Request, &1))).()
    request = case storage do
      :zookeeper -> %{request | api_version: 0}
      :kafka     -> %{request | api_version: 1}
    end

    # TODO
    # Handle every partition errors
    case Connection.request(conn, request) do
      {:ok, %{topics: [{^topic, partitions}]}} -> {:ok, partitions}
      {:error, reason} ->
        {:error, reason}
    end
  end

  defp get_offset(topic, partition, conn, time) when is_integer(partition) do
    request = %OffsetRequest{topics: [{topic, [{partition, time, 1}]}]}
    case Connection.request(conn, request) do
      {:ok, %{offsets: [{_, [%{error: :no_error, offsets: [offset]}]}]}} ->
        {:ok, {offset, ""}}
      {:ok, %{offsets: [{_, [%{error: :no_error, offsets: []}]}]}} ->
        {:ok, {0, ""}}
      {:ok, %{offsets: [{_, [%{error: error}]}]}} ->
        {:error, error}
      error ->
        error
    end
  end

  defp cancel_timer(%{timer: :sent} = state), do: %{state | timer: nil}
  defp cancel_timer(%{timer: nil} = state), do: state
  defp cancel_timer(%{timer: timer} = state) do
    :erlang.cancel_timer(timer)
    %{state | timer: nil}
  end

  defp start_timer(%{timer: nil, interval: interval} = state) do
    timer = :erlang.start_timer(interval, self, :do_commit)
    %{state | timer: timer}
  end
  defp start_timer(state), do: state
end