% 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
%
%   http://www.apache.org/licenses/LICENSE-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.

-module(fabric_streams).

-export([
    start/2,
    start/3,
    start/4,
    start/5,
    cleanup/1
]).

-include_lib("fabric/include/fabric.hrl").
-include_lib("mem3/include/mem3.hrl").


-define(WORKER_CLEANER, fabric_worker_cleaner).


start(Workers, Keypos) ->
    start(Workers, Keypos, undefined, undefined).


start(Workers, Keypos, RingOpts) ->
    start(Workers, Keypos, undefined, undefined, RingOpts).


start(Workers, Keypos, StartFun, Replacements) ->
    start(Workers, Keypos, StartFun, Replacements, []).


start(Workers0, Keypos, StartFun, Replacements, RingOpts) ->
    Fun = fun handle_stream_start/3,
    Acc = #stream_acc{
        workers = fabric_dict:init(Workers0, waiting),
        ready = [],
        start_fun = StartFun,
        replacements = Replacements,
        ring_opts = RingOpts
    },
    spawn_worker_cleaner(self(), Workers0),
    Timeout = fabric_util:request_timeout(),
    case rexi_utils:recv(Workers0, Keypos, Fun, Acc, Timeout, infinity) of
        {ok, #stream_acc{ready = Workers}} ->
            AckedWorkers = fabric_dict:fold(fun(Worker, From, WorkerAcc) ->
                rexi:stream_start(From),
                [Worker | WorkerAcc]
            end, [], Workers),
            {ok, AckedWorkers};
        Else ->
            Else
    end.


cleanup(Workers) ->
    % Stop the auxiliary cleaner process as we got to the point where cleanup
    % happesn in the regular fashion so we don't want to send 2x the number kill
    % messages
    case get(?WORKER_CLEANER) of
        CleanerPid when is_pid(CleanerPid) ->
            erase(?WORKER_CLEANER),
            exit(CleanerPid, kill);
        _ ->
            ok
    end,
    fabric_util:cleanup(Workers).


handle_stream_start({rexi_DOWN, _, {_, NodeRef}, _}, _, St) ->
    #stream_acc{workers = Workers, ready = Ready, ring_opts = RingOpts} = St,
    case fabric_ring:node_down(NodeRef, Workers, Ready, RingOpts) of
        {ok, Workers1} ->
            {ok, St#stream_acc{workers = Workers1}};
        error ->
            {error, {nodedown, <<"progress not possible">>}}
    end;

handle_stream_start({rexi_EXIT, Reason}, Worker, St) ->
    #stream_acc{
        workers = Workers,
        ready = Ready,
        replacements = Replacements,
        ring_opts = RingOpts
    } = St,
    case {fabric_ring:handle_error(Worker, Workers, Ready, RingOpts), Reason} of
        {{ok, Workers1}, _Reason} ->
            {ok, St#stream_acc{workers = Workers1}};
        {error, {maintenance_mode, _Node}} when Replacements /= undefined ->
            % Check if we have replacements for this range
            % and start the new workers if so.
            case lists:keytake(Worker#shard.range, 1, Replacements) of
                {value, {_Range, WorkerReplacements}, NewReplacements} ->
                    FinalWorkers = lists:foldl(fun(Repl, NewWorkers) ->
                        NewWorker = (St#stream_acc.start_fun)(Repl),
                        add_worker_to_cleaner(self(), NewWorker),
                        fabric_dict:store(NewWorker, waiting, NewWorkers)
                    end, Workers, WorkerReplacements),
                    % Assert that our replaced worker provides us
                    % the oppurtunity to make progress. Need to make sure
                    % to include already processed responses, since we are
                    % checking the full range and some workers have already
                    % responded and were removed from the workers list
                    ReadyWorkers = [{W, R} || {_, W, R} <- Ready],
                    AllWorkers = FinalWorkers ++ ReadyWorkers,
                    true = fabric_ring:is_progress_possible(AllWorkers),
                    NewRefs = fabric_dict:fetch_keys(FinalWorkers),
                    {new_refs, NewRefs, St#stream_acc{
                        workers=FinalWorkers,
                        replacements=NewReplacements
                    }};
                false ->
                    % If we progress isn't possible and we don't have any
                    % replacements then we're dead in the water.
                    {error, {nodedown, <<"progress not possible">>}}
            end;
        {error, _} ->
            {error, fabric_util:error_info(Reason)}
    end;

handle_stream_start(rexi_STREAM_INIT, {Worker, From}, St) ->
    #stream_acc{workers = Workers, ready = Ready, ring_opts = RingOpts} = St,
    case fabric_dict:lookup_element(Worker, Workers) of
    undefined ->
        % This worker lost the race with other partition copies, terminate
        rexi:stream_cancel(From),
        {ok, St};
    waiting ->
        case fabric_ring:handle_response(Worker, From, Workers, Ready, RingOpts) of
            {ok, {Workers1, Ready1}} ->
                % Don't have a full ring yet. Keep getting responses
                {ok, St#stream_acc{workers = Workers1, ready = Ready1}};
            {stop, Ready1} ->
                % Have a full ring of workers. But don't ack the worker
                % yet so they don't start sending us rows until we're ready
                {stop, St#stream_acc{workers = [], ready = Ready1}}
        end
    end;

handle_stream_start({ok, ddoc_updated}, _, St) ->
    WaitingWorkers = [W || {W, _} <- St#stream_acc.workers],
    ReadyWorkers = [W || {W, _} <- St#stream_acc.ready],
    cleanup(WaitingWorkers ++ ReadyWorkers),
    {stop, ddoc_updated};

handle_stream_start(Else, _, _) ->
    exit({invalid_stream_start, Else}).


% Spawn an auxiliary rexi worker cleaner. This will be used in cases
% when the coordinator (request) process is forceably killed and doesn't
% get a chance to process its `after` fabric:clean/1 clause.
spawn_worker_cleaner(Coordinator, Workers) ->
    case get(?WORKER_CLEANER) of
        undefined ->
            Pid = spawn(fun() ->
                erlang:monitor(process, Coordinator),
                cleaner_loop(Coordinator, Workers)
            end),
            put(?WORKER_CLEANER, Pid),
            Pid;
         ExistingCleaner ->
            ExistingCleaner
   end.


cleaner_loop(Pid, Workers) ->
    receive
        {add_worker, Pid, Worker} ->
            cleaner_loop(Pid, [Worker | Workers]);
        {'DOWN', _, _, Pid, _} ->
            fabric_util:cleanup(Workers)
    end.


add_worker_to_cleaner(CoordinatorPid, Worker) ->
    case get(?WORKER_CLEANER) of
        CleanerPid when is_pid(CleanerPid) ->
            CleanerPid ! {add_worker, CoordinatorPid, Worker};
        _ ->
            ok
    end.


-ifdef(TEST).

-include_lib("eunit/include/eunit.hrl").

worker_cleaner_test_() ->
    {
        "Fabric spawn_worker_cleaner test", {
            setup, fun setup/0, fun teardown/1,
            fun(_) -> [
                should_clean_workers(),
                does_not_fire_if_cleanup_called(),
                should_clean_additional_worker_too()
            ] end
        }
    }.


should_clean_workers() ->
    ?_test(begin
        meck:reset(rexi),
        erase(?WORKER_CLEANER),
        Workers = [
            #shard{node = 'n1', ref = make_ref()},
            #shard{node = 'n2', ref = make_ref()}
        ],
        {Coord, _} = spawn_monitor(fun() -> receive die -> ok end end),
        Cleaner = spawn_worker_cleaner(Coord, Workers),
        Ref = erlang:monitor(process, Cleaner),
        Coord ! die,
        receive {'DOWN', Ref, _, Cleaner, _} -> ok end,
        ?assertEqual(1, meck:num_calls(rexi, kill_all, 1))
    end).


does_not_fire_if_cleanup_called() ->
    ?_test(begin
        meck:reset(rexi),
        erase(?WORKER_CLEANER),
        Workers = [
            #shard{node = 'n1', ref = make_ref()},
            #shard{node = 'n2', ref = make_ref()}
        ],
        {Coord, _} = spawn_monitor(fun() -> receive die -> ok end end),
        Cleaner = spawn_worker_cleaner(Coord, Workers),
        Ref = erlang:monitor(process, Cleaner),
        cleanup(Workers),
        Coord ! die,
        receive {'DOWN', Ref, _, _, _} -> ok end,
        % 2 calls would be from cleanup/1 function. If cleanup process fired
        % too it would have been 4 calls total.
        ?assertEqual(1, meck:num_calls(rexi, kill_all, 1))
    end).


should_clean_additional_worker_too() ->
    ?_test(begin
        meck:reset(rexi),
        erase(?WORKER_CLEANER),
        Workers = [
            #shard{node = 'n1', ref = make_ref()}
        ],
        {Coord, _} = spawn_monitor(fun() -> receive die -> ok end end),
        Cleaner = spawn_worker_cleaner(Coord, Workers),
        add_worker_to_cleaner(Coord, #shard{node = 'n2', ref = make_ref()}),
        Ref = erlang:monitor(process, Cleaner),
        Coord ! die,
        receive {'DOWN', Ref, _, Cleaner, _} -> ok end,
        ?assertEqual(1, meck:num_calls(rexi, kill_all, 1))
    end).


setup() ->
    ok = meck:expect(rexi, kill_all, fun(_) -> ok end).


teardown(_) ->
    meck:unload().

-endif.