%%%-------------------------------------------------------------------
%%% @author anlicheng
%%% @copyright (C) 2024, <COMPANY>
%%% @doc
%%%
%%% @end
%%% Created : 03. 9月 2024 11:32
%%%-------------------------------------------------------------------
-module(influx_client_pool).
-author("anlicheng").

-behaviour(gen_server).

%% API
-export([start_link/0]).
-export([write_data/4]).

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

-define(SERVER, ?MODULE).

%% buffer的数量
-define(BUFFER_SIZE, 1500).

-define(DEFAULT_BUCKET, <<"metric">>).
-define(DEFAULT_ORG, <<"nannong">>).

-record(state, {
    pool_pid :: pid(),
    %% 缓冲区, 格式: #{bucket => []}
    buffers = #{}
}).

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

-spec write_data(Measurement :: binary(), Tags :: map(), FieldsList :: list(), Timestamp :: integer()) -> no_return().
write_data(Measurement, Tags, FieldsList, Timestamp) when is_binary(Measurement), is_map(Tags), is_list(FieldsList), is_integer(Timestamp) ->
    %% 过来掉没有key的选项
    Points = lists:flatmap(fun(Fields) ->
        case Fields of
            #{<<"key">> := Key} ->
                Values = maps:remove(<<"key">>, Fields),
                NFields = #{Key => Values},
                Point = influx_point:new(Measurement, Tags, NFields, format_timestamp(Timestamp)),
                [Point];
            _ ->
                []
        end
                           end, FieldsList),
    Bucket = get_bucket(Measurement),

    gen_server:cast(?SERVER, {write_data, Bucket, Points}).

format_timestamp(Timestamp) when is_integer(Timestamp) ->
    case length(integer_to_list(Timestamp)) of
        10 -> Timestamp * 1000;
        13 -> Timestamp;
        16 -> Timestamp div 1000;
        19 -> Timestamp div 1000_000;
        _ -> Timestamp
    end.

-spec get_bucket(DeviceUUID :: binary()) -> binary().
get_bucket(DeviceUUID) when is_binary(DeviceUUID) ->
    SuperDeviceUUIDs = iot_env:get_value(super_device_uuids, []),
    case lists:member(DeviceUUID, SuperDeviceUUIDs) of
        true ->
            <<"metric_", DeviceUUID/binary>>;
        false ->
            ?DEFAULT_BUCKET
    end.

%% @doc Spawns the server and registers the local name (unique)
-spec(start_link() ->
    {ok, Pid :: pid()} | ignore | {error, Reason :: term()}).
start_link() ->
    gen_server:start_link({local, ?SERVER}, ?MODULE, [], []).

%%%===================================================================
%%% gen_server callbacks
%%%===================================================================

%% @private
%% @doc Initializes the server
-spec(init(Args :: term()) ->
    {ok, State :: #state{}} | {ok, State :: #state{}, timeout() | hibernate} |
    {stop, Reason :: term()} | ignore).
init([]) ->
    {ok, PoolProps} = application:get_env(iot, influx_pool),
    PoolSize = proplists:get_value(pool_size, PoolProps),
    WorkerArgs = proplists:get_value(worker_args, PoolProps),

    %% 启动工作的线程池
    {ok, PoolPid} = poolboy:start_link([{size, PoolSize}, {max_overflow, PoolSize}, {worker_module, influx_client}], WorkerArgs),
    %% 创建特殊的buckets
    SuperDeviceUUIDs = iot_env:get_value(super_device_uuids, []),
    Result = catch create_buckets(PoolPid, SuperDeviceUUIDs),
    lager:debug("[influx_client_pool] create buckets result is: ~p", [Result]),

    %% 定时刷新逻辑
    erlang:start_timer(5000, self(), flush_ticker),

    {ok, #state{pool_pid = PoolPid}}.

%% @private
%% @doc Handling call messages
-spec(handle_call(Request :: term(), From :: {pid(), Tag :: term()},
    State :: #state{}) ->
    {reply, Reply :: term(), NewState :: #state{}} |
    {reply, Reply :: term(), NewState :: #state{}, timeout() | hibernate} |
    {noreply, NewState :: #state{}} |
    {noreply, NewState :: #state{}, timeout() | hibernate} |
    {stop, Reason :: term(), Reply :: term(), NewState :: #state{}} |
    {stop, Reason :: term(), NewState :: #state{}}).
handle_call(_Request, _From, State = #state{}) ->
    {reply, ok, State}.

%% @private
%% @doc Handling cast messages
-spec(handle_cast(Request :: term(), State :: #state{}) ->
    {noreply, NewState :: #state{}} |
    {noreply, NewState :: #state{}, timeout() | hibernate} |
    {stop, Reason :: term(), NewState :: #state{}}).
handle_cast({write_data, _Bucket, []}, State) ->
    {noreply, State};
handle_cast({write_data, Bucket, Points}, State = #state{pool_pid = PoolPid, buffers = Buffers}) ->
    Q = maps:get(Bucket, Buffers, queue:new()),
    NQ = lists:foldl(fun(Point, Q0) -> queue:in(Point, Q0) end, Q, Points),

    %% 超过缓冲区设置的大小则批量导入
    Len = queue:len(NQ),
    case Len >= ?BUFFER_SIZE of
        true ->
            lager:debug("[influx_client_pool] bucket: ~p, flush buffer size: ~p", [Bucket, Len]),
            poolboy:transaction(PoolPid, fun(Pid) -> influx_client:write(Pid, Bucket, ?DEFAULT_ORG, <<"ms">>, queue:to_list(NQ)) end),
            {noreply, State#state{buffers = maps:put(Bucket, queue:new(), Buffers)}};
        false ->
            {noreply, State#state{buffers = maps:put(Bucket, NQ, Buffers)}}
    end.

%% @private
%% @doc Handling all non call/cast messages
-spec(handle_info(Info :: timeout() | term(), State :: #state{}) ->
    {noreply, NewState :: #state{}} |
    {noreply, NewState :: #state{}, timeout() | hibernate} |
    {stop, Reason :: term(), NewState :: #state{}}).
handle_info({timeout, _, flush_ticker}, State = #state{buffers = Buffers, pool_pid = PoolPid}) ->
    erlang:start_timer(5000, self(), flush_ticker),

    Acc = lists:foldl(fun({Bucket, Q}, Acc0) ->
        Len = queue:len(Q),
        case Len > 0 of
            true ->
                poolboy:transaction(PoolPid, fun(Pid) -> influx_client:write(Pid, Bucket, ?DEFAULT_ORG, <<"ms">>, queue:to_list(Q)) end),
                Acc0 + Len;
            false ->
                Acc0
        end
                      end, 0, maps:to_list(Buffers)),
    lager:debug("[influx_client_pool] flush_ticker acc write num: ~p", [Acc]),

    {noreply, State#state{buffers = #{}}};

handle_info(_Info, State = #state{}) ->
    {noreply, State}.

%% @private
%% @doc This function is called by a gen_server when it is about to
%% terminate. It should be the opposite of Module:init/1 and do any
%% necessary cleaning up. When it returns, the gen_server terminates
%% with Reason. The return value is ignored.
-spec(terminate(Reason :: (normal | shutdown | {shutdown, term()} | term()),
    State :: #state{}) -> term()).
terminate(_Reason, _State = #state{}) ->
    ok.

%% @private
%% @doc Convert process state when code is changed
-spec(code_change(OldVsn :: term() | {down, term()}, State :: #state{},
    Extra :: term()) ->
    {ok, NewState :: #state{}} | {error, Reason :: term()}).
code_change(_OldVsn, State = #state{}, _Extra) ->
    {ok, State}.

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

%% 创建特殊的buckets
-spec create_buckets(PoolPid :: pid(), SuperDeviceUUIDs :: list()) -> no_return().
create_buckets(PoolPid, []) when is_pid(PoolPid) ->
    ok;
create_buckets(PoolPid, SuperDeviceUUIDs) when is_pid(PoolPid), is_list(SuperDeviceUUIDs) ->
    poolboy:transaction(PoolPid, fun(Pid) ->
        {ok, Orgs} = influx_client:get_orgs(Pid),
        OrgMap = maps:from_list(lists:map(fun(#{<<"name">> := Name, <<"id">> := Id}) -> {Name, Id} end, Orgs)),
        OrgId = maps:get(?DEFAULT_ORG, OrgMap, <<"">>),
        lager:debug("[influx_client_pool] org_name: ~p, id is: ~p", [?DEFAULT_ORG, OrgId]),
        [influx_client:create_bucket(Pid, OrgId, <<"metric_", DeviceUUID/binary>>) || DeviceUUID <- SuperDeviceUUIDs]
                                 end).