%%%-------------------------------------------------------------------
%%% @author aresei
%%% @copyright (C) 2023, <COMPANY>
%%% @doc
%%%
%%% @end
%%% Created : 06. 7月 2023 12:02
%%%-------------------------------------------------------------------
-module(endpoint_mqtt).

-include("endpoint.hrl").
-behaviour(gen_statem).

%% API
-export([start_link/2]).

%% gen_statem callbacks
-export([callback_mode/0, init/1, terminate/3, code_change/4]).
-export([disconnected/3, connected/3]).

%% 消息重发间隔
-define(RETRY_INTERVAL, 15000).

-record(state, {
    endpoint :: #endpoint{},
    buffer :: endpoint_buffer:buffer(),
    conn_pid :: undefined | pid(),
    %% 待确认的数据, #{PacketId :: integer() => Id :: integer()}
    inflight = #{}
}).

-type mqtt_state() :: disconnected | connected.

%%%===================================================================
%%% API
%%%===================================================================

-spec start_link(LocalName :: term(), Endpoint :: #endpoint{}) ->
    {ok, pid()} | ignore | {error, term()}.
start_link(LocalName, Endpoint = #endpoint{}) ->
    gen_statem:start_link({via, gproc, {n, l, LocalName}}, ?MODULE, [Endpoint], []).

%%%===================================================================
%%% gen_statem callbacks
%%%===================================================================

-spec callback_mode() -> state_functions.
callback_mode() ->
    state_functions.

-spec init(term()) -> gen_statem:init_result(mqtt_state(), #state{}).
init([Endpoint = #endpoint{matcher = Matcher}]) ->
    ok = endpoint_util:set_metadata(),
    erlang:process_flag(trap_exit, true),
    ok = endpoint_subscription:subscribe(Matcher, self()),
    Buffer = endpoint_buffer:new(Endpoint, 10),
    {ok, disconnected, #state{endpoint = Endpoint, buffer = Buffer}, [{next_event, internal, connect}]}.

-spec disconnected(gen_statem:event_type(), term(), #state{}) ->
    gen_statem:event_handler_result(mqtt_state(), #state{}).
disconnected({call, From}, get_stat, State = #state{buffer = Buffer}) ->
    reply_stat(From, Buffer, State);
disconnected(cast, {forward, Metric}, State = #state{buffer = Buffer}) ->
    store_metric(Metric, Buffer, State);
disconnected(cast, cleanup, State = #state{buffer = Buffer}) ->
    cleanup_buffer(Buffer, State);
disconnected(cast, {reload, NEndpoint = #endpoint{matcher = NMatcher}},
             State = #state{endpoint = #endpoint{matcher = Matcher}, conn_pid = ConnPid}) ->
    reload_endpoint(Matcher, NMatcher, ConnPid, NEndpoint, State);
disconnected(internal, connect, State) ->
    {keep_state, State, [{next_event, internal, do_connect}]};
disconnected(state_timeout, connect, State) ->
    {keep_state, State, [{next_event, internal, do_connect}]};
disconnected(internal, do_connect,
             State = #state{
                 buffer = Buffer,
                 endpoint = #endpoint{
                     title = Title,
                     config = #mqtt_endpoint{
                         host = Host,
                         port = Port,
                         username = Username,
                         password = Password,
                         client_id = ClientId
                     }
                 }
             }) ->
    logger:debug("[endpoint_mqtt] endpoint: ~ts, create postman", [Title]),
    Opts = [
        {owner, self()},
        {clientid, ClientId},
        {host, binary_to_list(Host)},
        {port, Port},
        {tcp_opts, []},
        {username, binary_to_list(Username)},
        {password, binary_to_list(Password)},
        {keepalive, 86400},
        {auto_ack, true},
        {connect_timeout, 5000},
        {proto_ver, v5},
        {retry_interval, 5000}
    ],
    case connect_mqtt(Opts) of
        {ok, ConnPid} ->
            logger:debug("[endpoint_mqtt] connect success, pid: ~p", [ConnPid]),
            NBuffer = endpoint_buffer:trigger_n(Buffer),
            {next_state, connected, State#state{conn_pid = ConnPid, buffer = NBuffer}};
        {error, Reason} ->
            logger:warning("[endpoint_mqtt] connect get error: ~p", [Reason]),
            {keep_state, State#state{conn_pid = undefined, inflight = #{}},
             [{state_timeout, ?RETRY_INTERVAL, connect}]}
    end;
disconnected(info, {timeout, TimerRef, {endpoint_buffer_retry, Id, Payload}}, State = #state{buffer = Buffer}) ->
    NBuffer = endpoint_buffer:handle_timeout(TimerRef, Id, Payload, Buffer),
    {keep_state, State#state{buffer = NBuffer}};
disconnected(info, {next_data, _Id, _Tuple}, State) ->
    {keep_state, State};
disconnected(info, {'EXIT', ConnPid, Reason},
             State = #state{endpoint = #endpoint{title = Title}, conn_pid = ConnPid}) ->
    logger:warning("[endpoint_mqtt] endpoint: ~p, conn pid exit with reason: ~p", [Title, Reason]),
    {keep_state, State#state{conn_pid = undefined, inflight = #{}},
     [{state_timeout, ?RETRY_INTERVAL, connect}]};
disconnected(info, Info, State) ->
    unknown_info(Info, disconnected, State);
disconnected(EventType, EventContent, State) ->
    unknown_event(EventType, EventContent, disconnected, State).

-spec connected(gen_statem:event_type(), term(), #state{}) ->
    gen_statem:event_handler_result(mqtt_state(), #state{}).
connected({call, From}, get_stat, State = #state{buffer = Buffer}) ->
    reply_stat(From, Buffer, State);
connected(cast, {forward, Metric}, State = #state{buffer = Buffer}) ->
    forward_metric(Metric, Buffer, State);
connected(cast, cleanup, State = #state{buffer = Buffer}) ->
    cleanup_buffer(Buffer, State);
connected(cast, {reload, NEndpoint = #endpoint{matcher = NMatcher}},
          State = #state{endpoint = #endpoint{matcher = Matcher}, conn_pid = ConnPid}) ->
    reload_endpoint(Matcher, NMatcher, ConnPid, NEndpoint, State);
connected(internal, connect, State) ->
    {keep_state, State};
connected(info, {next_data, Id, Metric},
          State = #state{
              conn_pid = ConnPid,
              buffer = Buffer,
              inflight = InFlight,
              endpoint = #endpoint{config = #mqtt_endpoint{topic = Topic, qos = Qos}}
          }) ->
    logger:debug("[endpoint_mqtt] will publish topic: ~p, metric: ~p, qos: ~p", [Topic, Metric, Qos]),
    case emqtt:publish(ConnPid, Topic, #{}, Metric, [{qos, Qos}, {retain, true}]) of
        ok ->
            NBuffer = endpoint_buffer:ack(Id, Buffer),
            {keep_state, State#state{buffer = NBuffer}};
        {ok, PacketId} ->
            {keep_state, State#state{inflight = maps:put(PacketId, Id, InFlight)}};
        {error, Reason} ->
            logger:warning("[endpoint_mqtt] send message to topic: ~p, get error: ~p", [Topic, Reason]),
            stop_mqtt_conn(ConnPid),
            NBuffer = endpoint_buffer:recover_inflight(Buffer),
            {next_state, disconnected, State#state{conn_pid = undefined, inflight = #{}, buffer = NBuffer},
             [{state_timeout, ?RETRY_INTERVAL, connect}]}
    end;
connected(info, {timeout, TimerRef, {endpoint_buffer_retry, Id, Payload}}, State = #state{buffer = Buffer}) ->
    NBuffer = endpoint_buffer:handle_timeout(TimerRef, Id, Payload, Buffer),
    {keep_state, State#state{buffer = NBuffer}};
connected(info, {disconnected, ReasonCode, Properties},
          State = #state{conn_pid = ConnPid, buffer = Buffer}) ->
    logger:debug("[endpoint_mqtt] Recv a DISONNECT packet - ReasonCode: ~p, Properties: ~p", [ReasonCode, Properties]),
    stop_mqtt_conn(ConnPid),
    NBuffer = endpoint_buffer:recover_inflight(Buffer),
    {next_state, disconnected, State#state{conn_pid = undefined, inflight = #{}, buffer = NBuffer},
     [{state_timeout, ?RETRY_INTERVAL, connect}]};
connected(info, {publish, Message = #{packet_id := _PacketId, payload := Payload}}, State) ->
    logger:debug("[endpoint_mqtt] Recv a publish packet: ~p, payload: ~p", [Message, Payload]),
    {keep_state, State};
connected(info, {puback, #{packet_id := PacketId}},
          State = #state{inflight = Inflight, buffer = Buffer}) ->
    case maps:take(PacketId, Inflight) of
        {Id, RestInflight} ->
            NBuffer = endpoint_buffer:ack(Id, Buffer),
            {keep_state, State#state{buffer = NBuffer, inflight = RestInflight}};
        error ->
            {keep_state, State}
    end;
connected(info, {'EXIT', ConnPid, Reason},
          State = #state{endpoint = #endpoint{title = Title}, conn_pid = ConnPid, buffer = Buffer}) ->
    logger:warning("[endpoint_mqtt] endpoint: ~p, conn pid exit with reason: ~p", [Title, Reason]),
    NBuffer = endpoint_buffer:recover_inflight(Buffer),
    {next_state, disconnected, State#state{conn_pid = undefined, inflight = #{}, buffer = NBuffer},
     [{state_timeout, ?RETRY_INTERVAL, connect}]};
connected(info, Info, State) ->
    unknown_info(Info, connected, State);
connected(EventType, EventContent, State) ->
    unknown_event(EventType, EventContent, connected, State).

-spec terminate(term(), mqtt_state(), #state{}) -> term().
terminate(Reason, _StateName, #state{endpoint = #endpoint{title = Title}, buffer = Buffer, conn_pid = ConnPid}) ->
    logger:debug("[endpoint_mqtt] endpoint: ~p, terminate with reason: ~p", [Title, Reason]),
    stop_mqtt_conn(ConnPid),
    endpoint_buffer:cleanup(Buffer),
    ok.

-spec code_change(term() | {down, term()}, mqtt_state(), #state{}, term()) ->
    {ok, mqtt_state(), #state{}} | {error, term()}.
code_change(_OldVsn, StateName, State = #state{}, _Extra) ->
    {ok, StateName, State}.

%%%===================================================================
%%% Internal functions
%%%===================================================================

-spec reply_stat(gen_statem:from(), endpoint_buffer:buffer(), #state{}) ->
    gen_statem:event_handler_result(mqtt_state(), #state{}).
reply_stat(From, Buffer, State) ->
    Stat = endpoint_buffer:stat(Buffer),
    {keep_state, State, [{reply, From, {ok, Stat}}]}.

-spec forward_metric(binary(), endpoint_buffer:buffer(), #state{}) ->
    gen_statem:event_handler_result(mqtt_state(), #state{}).
forward_metric(Metric, Buffer, State) ->
    NBuffer = endpoint_buffer:append(Metric, Buffer),
    {keep_state, State#state{buffer = NBuffer}}.

-spec store_metric(binary(), endpoint_buffer:buffer(), #state{}) ->
    gen_statem:event_handler_result(mqtt_state(), #state{}).
store_metric(Metric, Buffer, State) ->
    NBuffer = endpoint_buffer:append_only(Metric, Buffer),
    {keep_state, State#state{buffer = NBuffer}}.

-spec cleanup_buffer(endpoint_buffer:buffer(), #state{}) ->
    gen_statem:event_handler_result(mqtt_state(), #state{}).
cleanup_buffer(Buffer, State) ->
    NBuffer = endpoint_buffer:cleanup(Buffer),
    {keep_state, State#state{buffer = NBuffer}}.

-spec reload_endpoint(binary(), binary(), undefined | pid(), #endpoint{}, #state{}) ->
    gen_statem:event_handler_result(mqtt_state(), #state{}).
reload_endpoint(Matcher, NMatcher, ConnPid, NEndpoint, State = #state{}) ->
    ensure_subscription(Matcher, NMatcher),
    stop_mqtt_conn(ConnPid),
    NBuffer = endpoint_buffer:recover_inflight(State#state.buffer),
    {next_state, disconnected, State#state{endpoint = NEndpoint, conn_pid = undefined, inflight = #{}, buffer = NBuffer},
     [{next_event, internal, do_connect}]}.

-spec unknown_info(term(), mqtt_state(), #state{}) ->
    gen_statem:event_handler_result(mqtt_state(), #state{}).
unknown_info(Info, StateName, State) ->
    logger:warning("[endpoint_mqtt] unknown message: ~p, status: ~p", [Info, StateName]),
    {keep_state, State}.

-spec unknown_event(gen_statem:event_type(), term(), mqtt_state(), #state{}) ->
    gen_statem:event_handler_result(mqtt_state(), #state{}).
unknown_event(EventType, EventContent, StateName, State) ->
    logger:warning("[endpoint_mqtt] unknown event: ~p, content: ~p, status: ~p", [EventType, EventContent, StateName]),
    {keep_state, State}.

-spec ensure_subscription(binary(), binary()) -> ok.
ensure_subscription(Matcher, Matcher) ->
    ok;
ensure_subscription(Matcher, NMatcher) ->
    ok = endpoint_subscription:unsubscribe(Matcher, self()),
    endpoint_subscription:subscribe(NMatcher, self()).

-spec connect_mqtt(list()) -> {ok, pid()} | {error, term()}.
connect_mqtt(Opts) ->
    try
        {ok, ConnPid} = emqtt:start_link(Opts),
        logger:debug("[endpoint_mqtt] start connect, options: ~p", [Opts]),
        case emqtt:connect(ConnPid, 5000) of
            {ok, _} ->
                {ok, ConnPid};
            {error, Reason} ->
                stop_mqtt_conn(ConnPid),
                {error, Reason}
        end
    catch
        _:Error ->
            {error, Error}
    end.

-spec stop_mqtt_conn(undefined | pid()) -> ok.
stop_mqtt_conn(undefined) ->
    ok;
stop_mqtt_conn(ConnPid) when is_pid(ConnPid) ->
    _ = catch emqtt:stop(ConnPid),
    ok.