%% Copyright (c) 2018-Present Pivotal Software, Inc. All Rights Reserved.
-module(gen_batch_server).

-export([start_link/3,
         start_link/4,
         init_it/6,
         stop/1,
         stop/3,
         cast/2,
         cast_batch/2,
         call/2,
         call/3,
         system_continue/3,
         system_terminate/4,
         system_get_state/1,
         write_debug/3,
         format_status/2
        ]).

-define(MIN_MAX_BATCH_SIZE, 32).
-define(MAX_MAX_BATCH_SIZE, 8192).

-type server_ref() :: pid() |
                      (LocalName :: atom()) |
                      {Name :: atom(), Node :: atom()} |
                      {'global', term()} |
                      {'via', Module :: module(), Name :: term()}.

-type from() :: {Pid :: pid(), Tag :: reference()}.

-type op() :: {cast, pid(), UserOp :: term()} |
              {call, from(), UserOp :: term()} |
              {info, UserOp :: term()}.

-record(config, {batch_size :: non_neg_integer(),
                 min_batch_size :: non_neg_integer(),
                 max_batch_size :: non_neg_integer(),
                 parent :: pid(),
                 name :: atom(),
                 module :: module()}).

-record(state, {batch = [] :: [op()],
                batch_count = 0 :: non_neg_integer(),
                config = #config{} :: #config{},
                state :: term(),
                debug :: list()}).

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

-export_type([from/0, op/0,
              action/0, server_ref/0]).

%%% Behaviour

-type action() ::
    {reply, from(), Msg :: term()}.
%% an action that can be returned from handle_batch/2

-callback init(Args :: term()) ->
    {ok, State} |
    {stop, Reason :: term()}
      when State :: term().

-callback handle_batch([op()], State) ->
    {ok, State} |
    {ok, [action()], State} |
    {stop, Reason :: term()}
      when State :: term().

-callback terminate(Reason :: term(), State :: term()) -> term().

-callback format_status(State :: term()) -> term().

%% TODO: code_change

-optional_callbacks([format_status/1,
                     terminate/2]).


%%%
%%% API
%%%

-spec start_link(Name, Mod, Args) -> Result when
     Name :: {local, atom()} | {global, term()} | {via, atom(), term()},
     Mod :: module(),
     Args :: term(),
     Result ::  {ok,pid()} | {error, {already_started, pid()}}.
start_link(Name, Mod, Args) ->
    gen:start(?MODULE, link, Name, Mod, Args, []).

-spec start_link(Name, Mod, Args, Options) -> Result when
     Name :: {local, atom()} | {global, term()} | {via, atom(), term()},
     Mod :: module(),
     Args :: term(),
     Options :: list(),
     Result ::  {ok,pid()} | {error, {already_started, pid()}}.
start_link(Name, Mod, Args, Opts) ->
    gen:start(?MODULE, link, Name, Mod, Args, Opts).

%% pretty much copied wholesale from gen_server
init_it(Starter, self, Name, Mod, Args, Options) ->
    init_it(Starter, self(), Name, Mod, Args, Options);
init_it(Starter, Parent, Name0, Mod, Args, Options) ->
    Name = gen:name(Name0),
    Debug = gen:debug_options(Name, Options),
    MaxBatchSize = proplists:get_value(max_batch_size, Options, ?MAX_MAX_BATCH_SIZE),
    MinBatchSize = proplists:get_value(min_batch_size, Options, ?MIN_MAX_BATCH_SIZE),
    case catch Mod:init(Args) of
        {ok, State0} ->
            proc_lib:init_ack(Starter, {ok, self()}),
            Conf = #config{module = Mod,
                           parent = Parent,
                           name = Name,
                           batch_size = MinBatchSize,
                           min_batch_size = MinBatchSize,
                           max_batch_size = MaxBatchSize},
            State = #state{config = Conf,
                           state = State0,
                           debug = Debug},
            loop_wait(State, Parent);
        {stop, Reason} ->
            %% For consistency, we must make sure that the
            %% registered name (if any) is unregistered before
            %% the parent process is notified about the failure.
            %% (Otherwise, the parent process could get
            %% an 'already_started' error if it immediately
            %% tried starting the process again.)
            gen:unregister_name(Name0),
            proc_lib:init_ack(Starter, {error, Reason}),
            exit(Reason);
        % ignore ->
        %     gen:unregister_name(Name0),
        %     proc_lib:init_ack(Starter, ignore),
        %     exit(normal);
        {'EXIT', Reason} ->
            gen:unregister_name(Name0),
            proc_lib:init_ack(Starter, {error, Reason}),
            exit(Reason);
        Else ->
            Error = {bad_return_value, Else},
            proc_lib:init_ack(Starter, {error, Error}),
            exit(Error)
    end.

stop(Name) ->
    gen:stop(Name).

stop(Name, Reason, Timeout) ->
    gen:stop(Name, Reason, Timeout).

-spec cast(server_ref(), term()) -> ok.
cast({global,Name}, Request) ->
    catch global:send(Name, cast_msg(Request)),
    ok;
cast({via, Mod, Name}, Request) ->
    catch Mod:send(Name, cast_msg(Request)),
    ok;
cast({Name,Node}=Dest, Request) when is_atom(Name), is_atom(Node) ->
    do_cast(Dest, Request);
cast(Dest, Request) when is_atom(Dest) ->
    do_cast(Dest, Request);
cast(Dest, Request) when is_pid(Dest) ->
    do_cast(Dest, Request).


do_cast(Dest, Request) ->
  do_send(Dest, cast_msg(Request)),
  ok.

cast_msg(Request) ->
  {'$gen_cast', Request}.

-spec cast_batch(server_ref(), [term()]) -> ok.
cast_batch({global,Name}, Batch) ->
    catch global:send(Name, cast_batch_msg(Batch)),
    ok;
cast_batch({via, Mod, Name}, Batch) ->
    catch Mod:send(Name, cast_batch_msg(Batch)),
    ok;
cast_batch({Name,Node}=Dest, Batch) when is_atom(Name), is_atom(Node) ->
    do_cast_batch(Dest, Batch);
cast_batch(Dest, Batch) when is_atom(Dest) ->
    do_cast_batch(Dest, Batch);
cast_batch(Dest, Batch) when is_pid(Dest) ->
    do_cast_batch(Dest, Batch).

do_cast_batch(Dest, Batch) ->
    do_send(Dest, cast_batch_msg(Batch)),
    ok.

cast_batch_msg(Msgs0) when is_list(Msgs0) ->
    Msgs = lists:foldl(fun (Msg, Acc) ->
                               [{cast, Msg} | Acc]
                       end, [], Msgs0),
    {'$gen_cast_batch', Msgs, length(Msgs)};
cast_batch_msg(_) ->
    erlang:error(badarg).

-spec call(server_ref(), Request :: term()) -> term().
call(Name, Request) ->
    call(Name, Request, 5000).

-spec call(pid() | atom(), term(), non_neg_integer()) -> term().
call(Name, Request, Timeout) ->
    case catch gen:call(Name, '$gen_call', Request, Timeout) of
        {ok, Res} ->
            Res;
        {'EXIT', Reason} ->
            exit({Reason, {?MODULE, call, [Name, Request, Timeout]}})
    end.


%% Internal

loop_wait(State0, Parent) ->
    %% batches can accumulate a lot of garbage, collect it here
    %% TOOD: make this configurable
    garbage_collect(),
    receive
        {system, From, Request} ->
            sys:handle_system_msg(Request, From, Parent,
                                  ?MODULE, State0#state.debug, State0);
        {'EXIT', Parent, Reason} ->
            terminate(Reason, State0),
            exit(Reason);
        Msg ->
            enter_loop_batched(Msg, Parent, State0)
    end.

append_msg({'$gen_cast', Msg},
           #state{batch = Batch,
                  batch_count = BatchCount} = State0) ->
    State0#state{batch = [{cast, Msg} | Batch],
                 batch_count = BatchCount + 1};
append_msg({'$gen_cast_batch', Msgs, Count},
           #state{batch = Batch,
                  batch_count = BatchCount} = State0) ->
    State0#state{batch = Msgs ++ Batch,
                 batch_count = BatchCount + Count};
append_msg({'$gen_call', From, Msg},
           #state{batch = Batch,
                  batch_count = BatchCount} = State0) ->
    State0#state{batch = [{call, From, Msg} | Batch],
                 batch_count = BatchCount + 1};
append_msg(Msg, #state{batch = Batch,
                       batch_count = BatchCount} = State0) ->
    State0#state{batch = [{info, Msg} | Batch],
                 batch_count = BatchCount + 1}.

enter_loop_batched(Msg, Parent, #state{debug = []} = State0) ->
    loop_batched(append_msg(Msg, State0), Parent);
enter_loop_batched(Msg, Parent, State0) ->
    State = handle_debug_in(State0, Msg),
    %% append to batch
    loop_batched(append_msg(Msg, State), Parent).

loop_batched(#state{config = #config{batch_size = BatchSize,
                                     max_batch_size = Max} = Config,
                    batch_count = BatchCount} = State0,
             Parent) when BatchCount >= BatchSize ->
    % complete batch after seeing batch_size writes
    State = complete_batch(State0),
    % grow max batch size
    NewBatchSize = min(Max, BatchSize * 2),
    loop_wait(State#state{config = Config#config{batch_size = NewBatchSize}},
              Parent);
loop_batched(#state{debug = Debug} = State0, Parent) ->
    receive
        {system, From, Request} ->
            sys:handle_system_msg(Request, From, Parent,
                                  ?MODULE, Debug, State0);
        {'EXIT', Parent, Reason} ->
            terminate(Reason, State0),
            exit(Reason);
        Msg ->
            enter_loop_batched(Msg, Parent, State0)
    after 0 ->
              State = complete_batch(State0),
              Config = State#state.config,
              NewBatchSize = max(Config#config.min_batch_size,
                                 Config#config.batch_size / 2),
              % loop_wait(State#state{batch_size = NewBatchSize}, Parent)
              loop_wait(State#state{config =
                                    Config#config{batch_size = NewBatchSize}},
                        Parent)
    end.

terminate(Reason, #state{config = #config{module = Mod}, state = Inner}) ->
    catch Mod:terminate(Reason, Inner),
    ok.


complete_batch(#state{batch = []} = State) ->
    State;
complete_batch(#state{batch = Batch,
                      config = #config{module = Mod},
                      state = Inner0,
                      debug = Debug0} = State0) ->
    case catch Mod:handle_batch(lists:reverse(Batch), Inner0) of
        {ok, Inner} ->
            State0#state{batch = [],
                         state = Inner,
                         batch_count = 0};
        {ok, Actions, Inner} ->
            Debug = lists:foldl(fun ({reply, {Pid, Tag}, Msg}, Dbg) ->
                                        Pid ! {Tag, Msg},
                                        handle_debug_out(Pid, Msg, Dbg)
                                end, Debug0, Actions),
            State0#state{batch = [],
                         state = Inner,
                         batch_count = 0,
                         debug = Debug};
        {stop, Reason} ->
            terminate(Reason, State0),
            exit(Reason);
        {'EXIT', Reason} ->
            terminate(Reason, State0),
            exit(Reason)
    end.

handle_debug_in(#state{debug = Dbg0} = State, Msg) ->
    Dbg = sys:handle_debug(Dbg0, fun write_debug/3, ?MODULE, {in, Msg}),
    State#state{debug = Dbg}.

handle_debug_out(_, _, []) ->
    [];
handle_debug_out(Pid, Msg, Dbg) ->
    Evt = {out, {self(), Msg}, Pid},
    sys:handle_debug(Dbg, fun write_debug/3, ?MODULE, Evt).

%% Here are the sys call back functions

system_continue(Parent, Debug, State) ->
    % TODO check if we've written to the current batch or not
    loop_batched(State#state{debug = Debug}, Parent).

-spec system_terminate(term(), pid(), list(), term()) -> no_return().
system_terminate(Reason, _Parent, _Debug, State) ->
    terminate(Reason, State),
    exit(Reason).

system_get_state(State) ->
    {ok, State}.

format_status(_Reason, [_PDict, SysState, Parent, Debug,
                        #state{config = #config{name = Name,
                                                module = Mod},
                               state = State }]) ->
    Header = gen:format_status_header("Status for batching server", Name),
    Log = sys:get_debug(log, Debug, []),
    [{header, Header},
     {data,
      [
       {"Status", SysState},
       {"Parent", Parent},
       {"Logged Events", Log}]} |
     case catch Mod:format_status(State) of
         L when is_list(L) -> L;
         {'EXIT', {undef, _}} ->
             %% not implemented just return the state
             [State];
         T -> [T]
     end].

write_debug(Dev, Event, Name) ->
    io:format(Dev, "~p event = ~p~n", [Name, Event]).

%% Send function

do_send(Dest, Msg) ->
    try erlang:send(Dest, Msg)
    catch
        error:_ -> ok
    end,
    ok.
