i3_client

-module(i3_client_connection).
-moduledoc false.

-export([callback_mode/0, init/1, handle_event/4]).
-export([
    start/3, start/4,
    start_link/3, start_link/4,
    start_monitor/3, start_monitor/4,
    call/2, call/3,
    reply/2
]).

-include("i3_client_codec.hrl").

-define(state, no_state).
-define(conn_opts, [path, auto_reconnect, reconnect_backoff, sync_connect]).
-define(badmsg(M), {badmsg, M}).

-type state() :: term().
-type command() :: {message_type(), term()}.
-type result() :: {message_type(), {ok, map()} | {error, enotsup | ?badmsg(map())}}.

-callback init(term()) -> {ok, state()}.
-callback publish(binary(), binary(), state()) -> ok.
-callback handle_connect(state()) -> {noreply, state()} | {send, term(), state()}.
-callback handle_disconnect(state()) -> {noreply, state()}.
-callback handle_call(command(), term(), state()) -> {noreply, state()} | {send, term(), state()}.
-callback handle_info(term(), state()) -> {noreply, state()} | {send, command(), state()}.
-callback handle_result(result(), state()) -> {noreply, state()}.

-optional_callbacks([
    handle_call/3,
    handle_connect/1,
    handle_disconnect/1,
    handle_info/2,
    handle_result/2
]).

-behaviour(gen_statem).

call(Server, Msg) ->
    call(Server, Msg, 5000).

call(Server, Msg, Timeout) ->
    maybe
        {?MODULE, Reason} ?= gen_statem:call(Server, {Msg, self()}, Timeout),
        exit({Reason, {?MODULE, call, [Server, Msg, Timeout]}})
    end.

reply({CallerPid, From}, Reply) when is_pid(CallerPid) ->
    gen_statem:reply(From, Reply).

start(Mod, Args, Opts) ->
    do_start(nolink, undefined, Mod, Args, Opts).

start(ServerName, Mod, Args, Opts) ->
    do_start(nolink, ServerName, Mod, Args, Opts).

start_link(Mod, Args, Opts) ->
    do_start(link, undefined, Mod, Args, Opts).

start_link(ServerName, Mod, Args, Opts) ->
    do_start(link, ServerName, Mod, Args, Opts).

start_monitor(Mod, Args, Opts) ->
    do_start(monitor, undefined, Mod, Args, Opts).

start_monitor(ServerName, Mod, Args, Opts) ->
    do_start(monitor, ServerName, Mod, Args, Opts).

do_start(LinkP, ServerName, Mod, Args, Opts) ->
    {ConnOpts0, StatemOpts} = i3_client_common:keystake(?conn_opts, Opts),
    ConnOpts1 =
        case lists:keyfind(sync_connect, 1, ConnOpts0) of
            false ->
                lists:keystore(sync_connect, 1, ConnOpts0, {sync_connect, true});
            _ ->
                ConnOpts0
        end,
    ConnOpts2 =
        case lists:keyfind(path, 1, ConnOpts1) of
            false ->
                Path = os:getenv("I3SOCK"),
                lists:keystore(path, 1, ConnOpts1, {path, Path});
            _ ->
                ConnOpts1
        end,
    case {LinkP, ServerName} of
        {nolink, undefined} ->
            gen_statem:start(?MODULE, {Mod, Args, ConnOpts2}, StatemOpts);
        {nolink, _} ->
            gen_statem:start(ServerName, ?MODULE, {Mod, Args, ConnOpts2}, StatemOpts);
        {link, undefined} ->
            gen_statem:start_link(?MODULE, {Mod, Args, ConnOpts2}, StatemOpts);
        {link, _} ->
            gen_statem:start_link(ServerName, ?MODULE, {Mod, Args, ConnOpts2}, StatemOpts);
        {monitor, undefined} ->
            gen_statem:start_monitor(?MODULE, {Mod, Args, ConnOpts2}, StatemOpts);
        {monitor, _} ->
            gen_statem:start_monitor(ServerName, ?MODULE, {Mod, Args, ConnOpts2}, StatemOpts)
    end.

callback_mode() ->
    handle_event_function.

init({Mod, Args, Opts0}) ->
    case Mod:init(Args) of
        {ok, ModState} ->
            {AutoReconnect, Opts1} = keypop(auto_reconnect, Opts0, false),
            {ReconnectBackoff, Opts2} = keypop(reconnect_backoff, Opts1, 500),
            State0 =
                #{
                    state => {Mod, ModState},
                    protocol => undefined,
                    version => undefined,
                    auto_reconnect => AutoReconnect,
                    reconnect_backoff => ReconnectBackoff,
                    connection_attempt => 0
                },
            erlang:put(Mod, Opts2),
            case proplists:get_bool(sync_connect, Opts2) of
                true ->
                    case handle_event(internal, {connect, init}, ?state, State0) of
                        {keep_state, State1} ->
                            {ok, ?state, State1};
                        {keep_state, State1, Actions} ->
                            {ok, ?state, State1, Actions};
                        {stop, Reason, _} ->
                            {stop, Reason}
                    end;
                false ->
                    {ok, ?state, State0, {next_event, internal, {connect, init}}}
            end
    end.

handle_event(internal, {connect, _}, ?state, #{state := {Mod, ModState}} = State0) ->
    Opts = erlang:get(Mod),
    State1 = maps:update_with(connection_attempt, fun(A) -> A + 1 end, State0),
    maybe
        {ok, Protocol0} ?= i3_client_protocol:connect(Opts),
        % In order to guard against unsupported messages we need to
        % obtain the version of the protocol. It includes the variant,
        % if any.
        {ok, {reply, ?get_version, Reply}, Protocol1} ?=
            i3_client_protocol:handle_message(?get_version, ?nopayload, Protocol0),
        {ok, Protocol2} ?= i3_client_protocol:checkin(Protocol1),
        {ok, Version} ?= i3_client_version:from_reply(Reply),
        maybe_handle(
            Mod,
            handle_connect,
            [ModState],
            State1#{protocol := Protocol2, version := Version}
        )
    else
        {error, Reason} ->
            case State1 of
                #{auto_reconnect := true, reconnect_backoff := Backoff} ->
                    {keep_state, State1, {
                        {timeout, backoff}, Backoff, undefined
                    }};
                _ ->
                    {stop, Reason, State1}
            end
    end;
handle_event({timeout, backoff}, undefined, ?state, State) ->
    {keep_state, State, {next_event, internal, {connect, reconnect}}};
handle_event(
    {call, From},
    {Msg, CallerPid},
    ?state,
    #{state := {Mod, ModState}, version := Version} = State
) ->
    CallbackFrom = {CallerPid, From},
    case i3_client_version:support(Version, Msg) of
        ok ->
            handle(Mod, handle_call, [Msg, CallbackFrom, ModState], From, State);
        {error, _} = Error ->
            % Reply directly because the callback module should not
            % care about unsupported messages.
            reply(CallbackFrom, Error),
            {keep_state, State}
    end;
handle_event(
    info,
    Msg,
    ?state,
    #{protocol := Protocol0, state := {Mod, ModState}} = State
) ->
    Publish = fun(Type, Payload) ->
        Mod:publish(Type, Payload, ModState)
    end,
    case i3_client_protocol:handle_info(Msg, Publish, Protocol0) of
        {ok, Protocol1} ->
            {keep_state, State#{protocol := Protocol1}};
        {unknown, Protocol1} ->
            maybe_handle(Mod, handle_info, [Msg, ModState], State#{protocol := Protocol1});
        {Error, Reason, Protocol1} ->
            reconnect_or_stop(Error, Reason, Protocol1, State)
    end.

maybe_handle(Mod, Fun, Args, State) ->
    case erlang:function_exported(Mod, Fun, erlang:length(Args)) of
        true ->
            handle(Mod, Fun, Args, undefined, State);
        false ->
            {keep_state, State}
    end.

handle(Mod, Fun, Args, From, State0) ->
    case erlang:apply(Mod, Fun, Args) of
        {noreply, ModState} ->
            {keep_state, State0#{state := {Mod, ModState}}};
        {send, Type, Payload, ModState} ->
            State1 = State0#{state := {Mod, ModState}},
            Publish = fun(Type1, Payload1) ->
                Mod:publish(Type1, Payload1, ModState)
            end,
            maybe
                {ok, Result, Protocol0} ?=
                    i3_client_protocol:handle_message(
                        Type,
                        Payload,
                        Publish,
                        maps:get(protocol, State1)
                    ),
                {ok, Protocol1} ?= i3_client_protocol:checkin(Protocol0),
                ResultTuple = normalise_result(Result),
                handle(Mod, handle_result, [ResultTuple, ModState], From, State1#{
                    protocol := Protocol1
                })
            else
                {disconnect, Reason, Protocol} ->
                    reconnect_or_stop(disconnect, Reason, Protocol, State1)
            end
    end.

reconnect_or_stop(Error, Reason, Protocol, #{state := {Mod, ModState}} = State0) when
    (Error =:= error) or (Error =:= disconnect)
->
    {keep_state, State1} = maybe_handle(Mod, handle_disconnect, [ModState], State0),
    case State1 of
        #{auto_reconnect := true} ->
            {keep_state, State1, {next_event, internal, {connect, reconnect}}};
        _ ->
            {stop, Reason, State1#{protocol := Protocol}}
    end.

keypop(Key, List0, Default) ->
    case lists:keytake(Key, 1, List0) of
        {value, {Key, Value}, List1} ->
            {Value, List1};
        false ->
            {Default, List0}
    end.

normalise_result({reply, Type, Msg0}) ->
    case json:decode(Msg0) of
        #{<<"success">> := false} = Msg1 ->
            {Type, {error, ?badmsg(Msg1)}};
        Msg1 when is_list(Msg1) ->
            IsFailure =
                fun
                    (#{<<"success">> := false}) -> true;
                    (_) -> false
                end,
            case lists:any(IsFailure, Msg1) of
                true -> {Type, {error, ?badmsg(Msg1)}};
                false -> {Type, {ok, Msg1}}
            end;
        Msg1 ->
            {Type, {ok, Msg1}}
    end.