%%%
    Copyright ( c ) 2017 - 2021 Klarna Bank AB ( publ )
%%%
    Licensed under the Apache License , Version 2.0 ( the " License " ) ;
%%%   you may not use this file except in compliance with the License.
%%%   You may obtain a copy of the License at
%%%
%%%       -2.0
%%%
%%%   Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an " AS IS " BASIS ,
%%%   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
%%%   See the License for the specific language governing permissions and
%%%   limitations under the License.
%%%

%%%=============================================================================
%%% @doc This is a utility module to help force commit offsets to kafka.
%%% @end
%%%=============================================================================

-module(brod_cg_commits).

-behaviour(gen_server).
-behaviour(brod_group_member).

-export([ run/2
        ]).

-export([ start_link/2
        , stop/1
        , sync/1
        ]).

%% callbacks for brod_group_coordinator
-export([ get_committed_offsets/2
        , assignments_received/4
        , assignments_revoked/1
        , assign_partitions/3
        ]).

-export([ code_change/3
        , handle_call/3
        , handle_cast/2
        , handle_info/2
        , init/1
        , terminate/2
        ]).

-include("brod_int.hrl").

-type topic() :: brod:topic().
-type partition() :: brod:partition().
-type offset() :: brod:offset().
-type group_id() :: brod:group_id().
-type member_id() :: brod:group_member_id().
-type retention() :: integer(). %% -1 to use whatever configured in kafka
-type offsets() :: latest | earliest | [{partition(), offset()}].
-type prop_key() :: id | topic | retention | protocol | offsets.
-type prop_val() :: group_id() | topic() | retention() | offsets()
                  | brod_group_coordinator:protocol_name().
-type group_input() :: [{prop_key(), prop_val()}].

-record(state,
        { client             :: brod:client()
        , groupId            :: brod:group_id()
        , memberId           :: ?undef | member_id()
        , generationId       :: ?undef | brod:group_generation_id()
        , coordinator        :: pid()
        , topic              :: ?undef | topic()
        , offsets            :: ?undef | offsets()
        , is_elected = false :: boolean()
        , pending_sync       :: ?undef | {pid(), reference()}
        , is_done = false    :: boolean()
        }).

-type state() :: #state{}.

%%%_* APIs =====================================================================

%% @doc Force commit offsets.
run(ClientId, GroupInput) ->
  {ok, Pid} = start_link(ClientId, GroupInput),
  ok = sync(Pid),
  ok = stop(Pid).

%% @doc Start (link) a group member.
%% The member will try to join the consumer group and
%% get assignments for the given topic-partitions,
%% then commit given offsets to kafka.
%% In case not all given partitions are assigned to it,
%% it will terminate with an exit exception
-spec start_link(brod:client(), group_input()) -> {ok, pid()} | {error, any()}.
start_link(Client, GroupInput) ->
  gen_server:start_link(?MODULE, {Client, GroupInput}, []).

%% @doc Stop the process.
-spec stop(pid()) -> ok.
stop(Pid) ->
  Mref = erlang:monitor(process, Pid),
  ok = gen_server:cast(Pid, stop),
  receive
    {'DOWN', Mref, process, Pid, _Reason} ->
      ok
  end.

%% @doc Make a call to the resetter process, the call will be blocked
%% until offsets are committed.
-spec sync(pid()) -> ok.
sync(Pid) ->
  ok = gen_server:call(Pid, sync, infinity).

%%%_* APIs for group coordinator ===============================================

%% @doc Called by group coordinator when there is new assignment received.
-spec assignments_received(pid(), member_id(), integer(),
                           brod:received_assignments()) -> ok.
assignments_received(Pid, MemberId, GenerationId, TopicAssignments) ->
  gen_server:cast(Pid, {new_assignments, MemberId,
                        GenerationId, TopicAssignments}).

%% @doc Called by group coordinator before re-joinning the consumer group.
-spec assignments_revoked(pid()) -> ok.
assignments_revoked(Pid) ->
  gen_server:call(Pid, unsubscribe_all_partitions, infinity).

%% @doc This function is called only when `partition_assignment_strategy'
%% is set for `callback_implemented' in group config.
-spec assign_partitions(pid(), [brod:group_member()],
                        [{brod:topic(), brod:partition()}]) ->
        [{member_id(), [brod:partition_assignment()]}].
assign_partitions(Pid, Members, TopicPartitionList) ->
  Call = {assign_partitions, Members, TopicPartitionList},
  gen_server:call(Pid, Call, infinity).

%% @doc Called by group coordinator when initializing the assignments
%% for subscriber.
%% NOTE: this function is called only when it is DISABLED to commit offsets
        to kafka . i.e.  is set to consumer_managed
-spec get_committed_offsets(pid(), [{brod:topic(), brod:partition()}]) ->
        {ok, [{{brod:topic(), brod:partition()}, brod:offset()}]}.
get_committed_offsets(_Pid, _TopicPartitions) -> {ok, []}.

%%%_* gen_server callbacks =====================================================

init({Client, GroupInput}) ->
  ok = brod_utils:assert_client(Client),
  GroupId = proplists:get_value(id, GroupInput),
  ok = brod_utils:assert_group_id(GroupId),
  Topic = proplists:get_value(topic, GroupInput),
  ProtocolName = proplists:get_value(protocol, GroupInput),
  Retention = proplists:get_value(retention, GroupInput),
  Offsets = proplists:get_value(offsets, GroupInput),
  %% use callback_implemented strategy so I know I am elected leader
  %% when `assign_partitions' callback is called.
  Config = [ {partition_assignment_strategy, callback_implemented}
           , {offset_retention_seconds, Retention}
           , {protocol_name, ProtocolName}
           , {rejoin_delay_seconds, 2}
           ],
  {ok, Pid} = brod_group_coordinator:start_link(Client, GroupId, [Topic],
                                                Config, ?MODULE, self()),
  State = #state{ client      = Client
                , groupId     = GroupId
                , coordinator = Pid
                , topic       = Topic
                , offsets     = Offsets
                },
  {ok, State}.

handle_info(Info, State) ->
  log(State, info, "Info discarded:~p", [Info]),
  {noreply, State}.

handle_call(sync, From, State0) ->
  State1 = State0#state{pending_sync = From},
  State = maybe_reply_sync(State1),
  {noreply, State};
handle_call({assign_partitions, Members, TopicPartitions}, _From,
            #state{topic = MyTopic,
                   offsets = Offsets} = State) ->
  log(State, info, "Assigning all topic partitions to self", []),
  MyTP = [{MyTopic, P} || {P, _} <- Offsets],
  %% Assert that my topic partitions are included in
  %% subscriptions collected from ALL members
  Pred = fun(TP) -> not lists:member(TP, TopicPartitions) end,
  case lists:filter(Pred, MyTP) of
    [] ->
      %% all of the give partitions are valid
      ok;
    BadPartitions ->
      PartitionNumbers = [P || {_T, P} <- BadPartitions],
      log(State, error, "Nonexisting partitions in input: ~p",
          [PartitionNumbers]),
      erlang:exit({non_existing_partitions, PartitionNumbers})
  end,
  %% To keep it simple, assign all my topic-partitions to self
  %% but discard all other topic-partitions.
  %% After all, I will leave group as soon as offsets are committed
  Result = assign_all_to_self(Members, MyTP),
  {reply, Result, State#state{is_elected = true}};
handle_call(unsubscribe_all_partitions, _From, #state{} = State) ->
  %% nothing to do, because I do not subscribe to any partition
  {reply, ok, State};
handle_call(Call, _From, State) ->
  {reply, {error, {unknown_call, Call}}, State}.

handle_cast({new_assignments, _MemberId, GenerationId, Assignments},
            #state{ is_elected = IsLeader
                  , offsets = OffsetsToCommit
                  , coordinator = Pid
                  , topic =  MyTopic
                  } = State) ->
  %% Write a log if I am not a leader,
  %% hope the desired partitions are all assigned to me
  IsLeader orelse log(State, info, "Not elected", []),
  Groupped0 =
    brod_utils:group_per_key(
      fun(#brod_received_assignment{ topic        = Topic
                                   , partition    = Partition
                                   , begin_offset = Offset
                                   }) ->
          {Topic, {Partition, Offset}}
      end, Assignments),
  %% Discard other topics if for whatever reason the group leader assigns
  %% irrelevant topic-partitions to me
  Groupped = lists:filter(fun({Topic, _}) -> Topic =:= MyTopic end, Groupped0),
  log(State, info, "current offsets:\n~p", [Groupped]),
  %% Assert all desired partitions are in assignment
  case Groupped of
    [] ->
      log(State, error, "Topic ~s is not received in assignment", [MyTopic]),
      erlang:exit({bad_topic_assignment, Groupped0});
    [{MyTopic, PartitionOffsetList}] ->
      MyPartitions = [P || {P, _O} <- OffsetsToCommit],
      ReceivedPartitions = [P || {P, _O} <- PartitionOffsetList],
      case MyPartitions -- ReceivedPartitions of
        [] ->
          ok;
        Left ->
          log(State, error,
              "Partitions ~p are not received in assignment, "
              "There is probably another active group member subscribing "
              "to topic ~s, stop it and retry\n", [MyTopic, Left]),
          erlang:exit({unexpected_assignments, Left})
      end
  end,
  %% Stage all offsets in coordinator process
  lists:foreach(
    fun({Partition, Offset}) ->
        %% -1 here, because brod_group_coordinator +1 to commit
        OffsetToCommit = Offset - 1,
        brod_group_coordinator:ack(Pid, GenerationId, MyTopic,
                                   Partition, OffsetToCommit)
    end, OffsetsToCommit),
  %% Now force it to commit
  case brod_group_coordinator:commit_offsets(Pid) of
    ok -> ok;
    {error, Reason} ->
      log(State, error, "Failed to commit, reason:\n~p", [Reason]),
      erlang:exit(commit_failed)
  end,
  {noreply, set_done(State)};
handle_cast(stop, State) ->
  {stop, normal, State};
handle_cast(_Cast, State) ->
  {noreply, State}.

code_change(_OldVsn, State, _Extra) ->
  {ok, State}.

terminate(_Reason, #state{}) ->
  ok.

%%%_* Internal Functions =======================================================

set_done(State) ->
  maybe_reply_sync(State#state{is_done = true}).

maybe_reply_sync(#state{is_done = false} = State) ->
  State;
maybe_reply_sync(#state{pending_sync = ?undef} = State) ->
  State;
maybe_reply_sync(#state{pending_sync = From} = State) ->
  gen_server:reply(From, ok),
  log(State, info, "done\n", []),
  State#state{pending_sync = ?undef}.

%% I am the current leader because I am assigning partitions.
%% My member ID should be positioned at the head of the member list.
-spec assign_all_to_self([brod:group_member()], [{topic(), partition()}]) ->
        [{member_id(), [brod:partition_assignment()]}].
assign_all_to_self([{MyMemberId, _} | Members], TopicPartitions) ->
  Groupped = brod_utils:group_per_key(TopicPartitions),
  [ {MyMemberId, Groupped}
  | [{Id, []} || {Id, _MemberMeta} <- Members]
  ].

-spec log(state(), logger:level(), io:format(), [term()]) ->
  ok.
log(#state{groupId  = GroupId}, Level, Fmt, Args) ->
  ?BROD_LOG(Level,
            "Group member (~s,coor=~p):\n" ++ Fmt,
            [GroupId, self() | Args]).

%%%_* Emacs ====================================================================
%%% Local Variables:
%%% allout-layout: t
  erlang - indent - level : 2
%%% End:



   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       -2.0

   Unless required by applicable law or agreed to in writing, software
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

=============================================================================
 @doc This is a utility module to help force commit offsets to kafka.
 @end
=============================================================================
 callbacks for brod_group_coordinator
 -1 to use whatever configured in kafka
_* APIs =====================================================================
 @doc Force commit offsets.
 @doc Start (link) a group member.
 The member will try to join the consumer group and
 get assignments for the given topic-partitions,
 then commit given offsets to kafka.
 In case not all given partitions are assigned to it,
 it will terminate with an exit exception
 @doc Stop the process.
 @doc Make a call to the resetter process, the call will be blocked
 until offsets are committed.
_* APIs for group coordinator ===============================================
 @doc Called by group coordinator when there is new assignment received.
 @doc Called by group coordinator before re-joinning the consumer group.
 @doc This function is called only when `partition_assignment_strategy'
 is set for `callback_implemented' in group config.
 @doc Called by group coordinator when initializing the assignments
 for subscriber.
 NOTE: this function is called only when it is DISABLED to commit offsets
_* gen_server callbacks =====================================================
 use callback_implemented strategy so I know I am elected leader
 when `assign_partitions' callback is called.
 Assert that my topic partitions are included in
 subscriptions collected from ALL members
 all of the give partitions are valid
 To keep it simple, assign all my topic-partitions to self
 but discard all other topic-partitions.
 After all, I will leave group as soon as offsets are committed
 nothing to do, because I do not subscribe to any partition
 Write a log if I am not a leader,
 hope the desired partitions are all assigned to me
 Discard other topics if for whatever reason the group leader assigns
 irrelevant topic-partitions to me
 Assert all desired partitions are in assignment
 Stage all offsets in coordinator process
 -1 here, because brod_group_coordinator +1 to commit
 Now force it to commit
_* Internal Functions =======================================================
 I am the current leader because I am assigning partitions.
 My member ID should be positioned at the head of the member list.
_* Emacs ====================================================================
 Local Variables:
 allout-layout: t
 End:
    Copyright ( c ) 2017 - 2021 Klarna Bank AB ( publ )
    Licensed under the Apache License , Version 2.0 ( the " License " ) ;
    distributed under the License is distributed on an " AS IS " BASIS ,


-module(brod_cg_commits).

-behaviour(gen_server).
-behaviour(brod_group_member).

-export([ run/2
        ]).

-export([ start_link/2
        , stop/1
        , sync/1
        ]).

-export([ get_committed_offsets/2
        , assignments_received/4
        , assignments_revoked/1
        , assign_partitions/3
        ]).

-export([ code_change/3
        , handle_call/3
        , handle_cast/2
        , handle_info/2
        , init/1
        , terminate/2
        ]).

-include("brod_int.hrl").

-type topic() :: brod:topic().
-type partition() :: brod:partition().
-type offset() :: brod:offset().
-type group_id() :: brod:group_id().
-type member_id() :: brod:group_member_id().
-type offsets() :: latest | earliest | [{partition(), offset()}].
-type prop_key() :: id | topic | retention | protocol | offsets.
-type prop_val() :: group_id() | topic() | retention() | offsets()
                  | brod_group_coordinator:protocol_name().
-type group_input() :: [{prop_key(), prop_val()}].

-record(state,
        { client             :: brod:client()
        , groupId            :: brod:group_id()
        , memberId           :: ?undef | member_id()
        , generationId       :: ?undef | brod:group_generation_id()
        , coordinator        :: pid()
        , topic              :: ?undef | topic()
        , offsets            :: ?undef | offsets()
        , is_elected = false :: boolean()
        , pending_sync       :: ?undef | {pid(), reference()}
        , is_done = false    :: boolean()
        }).

-type state() :: #state{}.


run(ClientId, GroupInput) ->
  {ok, Pid} = start_link(ClientId, GroupInput),
  ok = sync(Pid),
  ok = stop(Pid).

-spec start_link(brod:client(), group_input()) -> {ok, pid()} | {error, any()}.
start_link(Client, GroupInput) ->
  gen_server:start_link(?MODULE, {Client, GroupInput}, []).

-spec stop(pid()) -> ok.
stop(Pid) ->
  Mref = erlang:monitor(process, Pid),
  ok = gen_server:cast(Pid, stop),
  receive
    {'DOWN', Mref, process, Pid, _Reason} ->
      ok
  end.

-spec sync(pid()) -> ok.
sync(Pid) ->
  ok = gen_server:call(Pid, sync, infinity).


-spec assignments_received(pid(), member_id(), integer(),
                           brod:received_assignments()) -> ok.
assignments_received(Pid, MemberId, GenerationId, TopicAssignments) ->
  gen_server:cast(Pid, {new_assignments, MemberId,
                        GenerationId, TopicAssignments}).

-spec assignments_revoked(pid()) -> ok.
assignments_revoked(Pid) ->
  gen_server:call(Pid, unsubscribe_all_partitions, infinity).

-spec assign_partitions(pid(), [brod:group_member()],
                        [{brod:topic(), brod:partition()}]) ->
        [{member_id(), [brod:partition_assignment()]}].
assign_partitions(Pid, Members, TopicPartitionList) ->
  Call = {assign_partitions, Members, TopicPartitionList},
  gen_server:call(Pid, Call, infinity).

        to kafka . i.e.  is set to consumer_managed
-spec get_committed_offsets(pid(), [{brod:topic(), brod:partition()}]) ->
        {ok, [{{brod:topic(), brod:partition()}, brod:offset()}]}.
get_committed_offsets(_Pid, _TopicPartitions) -> {ok, []}.


init({Client, GroupInput}) ->
  ok = brod_utils:assert_client(Client),
  GroupId = proplists:get_value(id, GroupInput),
  ok = brod_utils:assert_group_id(GroupId),
  Topic = proplists:get_value(topic, GroupInput),
  ProtocolName = proplists:get_value(protocol, GroupInput),
  Retention = proplists:get_value(retention, GroupInput),
  Offsets = proplists:get_value(offsets, GroupInput),
  Config = [ {partition_assignment_strategy, callback_implemented}
           , {offset_retention_seconds, Retention}
           , {protocol_name, ProtocolName}
           , {rejoin_delay_seconds, 2}
           ],
  {ok, Pid} = brod_group_coordinator:start_link(Client, GroupId, [Topic],
                                                Config, ?MODULE, self()),
  State = #state{ client      = Client
                , groupId     = GroupId
                , coordinator = Pid
                , topic       = Topic
                , offsets     = Offsets
                },
  {ok, State}.

handle_info(Info, State) ->
  log(State, info, "Info discarded:~p", [Info]),
  {noreply, State}.

handle_call(sync, From, State0) ->
  State1 = State0#state{pending_sync = From},
  State = maybe_reply_sync(State1),
  {noreply, State};
handle_call({assign_partitions, Members, TopicPartitions}, _From,
            #state{topic = MyTopic,
                   offsets = Offsets} = State) ->
  log(State, info, "Assigning all topic partitions to self", []),
  MyTP = [{MyTopic, P} || {P, _} <- Offsets],
  Pred = fun(TP) -> not lists:member(TP, TopicPartitions) end,
  case lists:filter(Pred, MyTP) of
    [] ->
      ok;
    BadPartitions ->
      PartitionNumbers = [P || {_T, P} <- BadPartitions],
      log(State, error, "Nonexisting partitions in input: ~p",
          [PartitionNumbers]),
      erlang:exit({non_existing_partitions, PartitionNumbers})
  end,
  Result = assign_all_to_self(Members, MyTP),
  {reply, Result, State#state{is_elected = true}};
handle_call(unsubscribe_all_partitions, _From, #state{} = State) ->
  {reply, ok, State};
handle_call(Call, _From, State) ->
  {reply, {error, {unknown_call, Call}}, State}.

handle_cast({new_assignments, _MemberId, GenerationId, Assignments},
            #state{ is_elected = IsLeader
                  , offsets = OffsetsToCommit
                  , coordinator = Pid
                  , topic =  MyTopic
                  } = State) ->
  IsLeader orelse log(State, info, "Not elected", []),
  Groupped0 =
    brod_utils:group_per_key(
      fun(#brod_received_assignment{ topic        = Topic
                                   , partition    = Partition
                                   , begin_offset = Offset
                                   }) ->
          {Topic, {Partition, Offset}}
      end, Assignments),
  Groupped = lists:filter(fun({Topic, _}) -> Topic =:= MyTopic end, Groupped0),
  log(State, info, "current offsets:\n~p", [Groupped]),
  case Groupped of
    [] ->
      log(State, error, "Topic ~s is not received in assignment", [MyTopic]),
      erlang:exit({bad_topic_assignment, Groupped0});
    [{MyTopic, PartitionOffsetList}] ->
      MyPartitions = [P || {P, _O} <- OffsetsToCommit],
      ReceivedPartitions = [P || {P, _O} <- PartitionOffsetList],
      case MyPartitions -- ReceivedPartitions of
        [] ->
          ok;
        Left ->
          log(State, error,
              "Partitions ~p are not received in assignment, "
              "There is probably another active group member subscribing "
              "to topic ~s, stop it and retry\n", [MyTopic, Left]),
          erlang:exit({unexpected_assignments, Left})
      end
  end,
  lists:foreach(
    fun({Partition, Offset}) ->
        OffsetToCommit = Offset - 1,
        brod_group_coordinator:ack(Pid, GenerationId, MyTopic,
                                   Partition, OffsetToCommit)
    end, OffsetsToCommit),
  case brod_group_coordinator:commit_offsets(Pid) of
    ok -> ok;
    {error, Reason} ->
      log(State, error, "Failed to commit, reason:\n~p", [Reason]),
      erlang:exit(commit_failed)
  end,
  {noreply, set_done(State)};
handle_cast(stop, State) ->
  {stop, normal, State};
handle_cast(_Cast, State) ->
  {noreply, State}.

code_change(_OldVsn, State, _Extra) ->
  {ok, State}.

terminate(_Reason, #state{}) ->
  ok.


set_done(State) ->
  maybe_reply_sync(State#state{is_done = true}).

maybe_reply_sync(#state{is_done = false} = State) ->
  State;
maybe_reply_sync(#state{pending_sync = ?undef} = State) ->
  State;
maybe_reply_sync(#state{pending_sync = From} = State) ->
  gen_server:reply(From, ok),
  log(State, info, "done\n", []),
  State#state{pending_sync = ?undef}.

-spec assign_all_to_self([brod:group_member()], [{topic(), partition()}]) ->
        [{member_id(), [brod:partition_assignment()]}].
assign_all_to_self([{MyMemberId, _} | Members], TopicPartitions) ->
  Groupped = brod_utils:group_per_key(TopicPartitions),
  [ {MyMemberId, Groupped}
  | [{Id, []} || {Id, _MemberMeta} <- Members]
  ].

-spec log(state(), logger:level(), io:format(), [term()]) ->
  ok.
log(#state{groupId  = GroupId}, Level, Fmt, Args) ->
  ?BROD_LOG(Level,
            "Group member (~s,coor=~p):\n" ++ Fmt,
            [GroupId, self() | Args]).

  erlang - indent - level : 2
