%%%-------------------------------------------------------------------
%%% @author anlicheng
%%% @copyright (C) 2025, <COMPANY>
%%% @doc
%%% Endpoint 数据分发订阅索引。
%%%
%%% 当前实现面向较大订阅量场景，publish 热路径不再全量遍历订阅表，而是把
%%% 订阅拆成 exact 和 wildcard 两类：
%%%
%%% <ul>
%%%   <li>精确订阅写入 endpoint_subscription_exact，publish 时按 RouteKey
%%%       直接 ets:lookup/2。</li>
%%%   <li>包含通配符的订阅写入 ETS trie。trie 边存放在
%%%       endpoint_subscription_trie_edge，节点订阅存放在
%%%       endpoint_subscription_trie_sub。</li>
%%%   <li>endpoint_subscription_reverse 保存 SubscriberPid 到订阅对象的反向索引，
%%%       用于 unsubscribe 和进程 DOWN 清理。</li>
%%%   <li>endpoint_subscription_pid 记录每个 SubscriberPid 的 monitor 和订阅计数，
%%%       同一个 pid 只 monitor 一次。</li>
%%% </ul>
%%%
%%% 所有 matcher 共用同一个 root node，root node id 固定为 0。新 trie 节点使用
%%% 单调递增的 integer node id。subscribe/unsubscribe 通过本 gen_server 串行写 ETS；
%%% publish/2 直接并发读取 ETS，避免中心 gen_server 成为数据分发瓶颈。
%%%
%%% 匹配规则：
%%%
%%% <ul>
%%%   <li>普通 segment 必须完全相等，例如 <<"device/a/temp">> 只精确匹配
%%%       同名订阅。</li>
%%%   <li><<"*">> 是单级通配，只匹配一个 segment。例如
%%%       <<"device/*/temp">> 匹配 <<"device/a/temp">>，不匹配
%%%       <<"device/a/b/temp">>。</li>
%%%   <li><<"+">> 是末尾多级通配，只允许出现在 matcher 最后一段，并且至少
%%%       匹配一个剩余 segment。例如 <<"device/+">> 匹配 <<"device/a">>
%%%       和 <<"device/a/temp">>，不匹配 <<"device">>。</li>
%%% </ul>
%%%
%%% publish/2 的匹配流程：
%%%
%%% <ol>
%%%   <li>从 exact 表直接查找 RouteKey 对应的订阅。</li>
%%%   <li>把 RouteKey 按 "/" 拆分后，从 trie root 开始逐层匹配。每层最多查找
%%%       当前 segment 分支和 <<"*">> 分支。</li>
%%%   <li>每消费一层前，收集当前候选节点上的 plus 订阅；全部 segment 消费完后，
%%%       收集候选节点上的 exact 订阅。</li>
%%%   <li>合并 exact 和 trie 命中结果，按原有优先级排序，再按 SubscriberPid 去重。
%%%       同一个 pid 同时被多个 matcher 命中时，保留优先级最高的订阅。</li>
%%%   <li>对最终订阅者执行 endpoint:forward/2。</li>
%%% </ol>
%%% @end
%%% Created : 07. 11月 2025 16:27
%%%-------------------------------------------------------------------
-module(endpoint_subscription).
-author("anlicheng").

-behaviour(gen_server).

%% API
-export([start_link/0]).
-export([subscribe/2, unsubscribe/2, publish/2, get_subscribers/0]).
-export([is_valid_components/1, of_components/1]).

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

-define(SERVER, ?MODULE).
-define(EXACT_TAB, endpoint_subscription_exact).
-define(TRIE_EDGE_TAB, endpoint_subscription_trie_edge).
-define(TRIE_SUB_TAB, endpoint_subscription_trie_sub).
-define(REVERSE_TAB, endpoint_subscription_reverse).
-define(PID_TAB, endpoint_subscription_pid).
-define(ROOT_NODE, 0).

%% 定义订阅者
-record(subscriber, {
    topic :: binary(),
    subscriber_pid :: pid(),
    components = [],
    monitor_ref :: undefined | reference(),
    %% 优先级
    %% 1. 完全匹配的topic优先级别最高
    %% 2. 带 * 的订阅
    %% 3. 带 + 的订阅
    order :: integer()
}).

-record(state, {
    exact_tid :: ets:tid(),
    edge_tid :: ets:tid(),
    trie_sub_tid :: ets:tid(),
    reverse_tid :: ets:tid(),
    pid_tid :: ets:tid(),
    next_node_id = 1 :: pos_integer()
}).

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

-spec subscribe(Topic :: binary(), SubscriberPid :: pid()) -> ok | {error, Reason :: binary()}.
subscribe(Topic, SubscriberPid) when is_binary(Topic), is_pid(SubscriberPid) ->
    gen_server:call(?SERVER, {subscribe, Topic, SubscriberPid}).

-spec unsubscribe(Topic :: binary(), SubscriberPid :: pid()) -> ok.
unsubscribe(Topic, SubscriberPid) when is_binary(Topic), is_pid(SubscriberPid) ->
    gen_server:call(?SERVER, {unsubscribe, Topic, SubscriberPid}).

-spec get_subscribers() -> {ok, Subscribers :: list()}.
get_subscribers() ->
    gen_server:call(?SERVER, get_subscribers).

-spec publish(RouteKey :: binary(), Content :: binary()) -> ok.
publish(RouteKey, Content) when is_binary(RouteKey), is_binary(Content) ->
    case ets:info(?EXACT_TAB) of
        undefined ->
            ok;
        _ ->
            MatchedSubscribers = match_route_key(RouteKey),
            lists:foreach(fun(#subscriber{subscriber_pid = SubscriberPid}) ->
                endpoint:forward(SubscriberPid, Content)
            end, MatchedSubscribers),
            maybe_log_unmatched_publish(RouteKey, Content, MatchedSubscribers),
            logger:debug("[endpoint_subscription] route_key: ~p, match_count: ~p", [RouteKey, length(MatchedSubscribers)]),
            ok
    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 = endpoint_util:set_metadata(),
    ExactTid = ets:new(?EXACT_TAB, [named_table, protected, bag, {read_concurrency, true}]),
    EdgeTid = ets:new(?TRIE_EDGE_TAB, [named_table, protected, set, {read_concurrency, true}]),
    TrieSubTid = ets:new(?TRIE_SUB_TAB, [named_table, protected, bag, {read_concurrency, true}]),
    ReverseTid = ets:new(?REVERSE_TAB, [named_table, protected, bag]),
    PidTid = ets:new(?PID_TAB, [named_table, protected, set]),
    {ok, #state{
        exact_tid = ExactTid,
        edge_tid = EdgeTid,
        trie_sub_tid = TrieSubTid,
        reverse_tid = ReverseTid,
        pid_tid = PidTid
    }}.

%% @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{}}).
%%  同一个SubscriberPid只能订阅同一个topic一次
handle_call(get_subscribers, _From, State = #state{exact_tid = ExactTid, trie_sub_tid = TrieSubTid}) ->
    Subscribers = exact_subscribers(ExactTid) ++ trie_subscribers(TrieSubTid),
    {reply, {ok, Subscribers}, State};

handle_call({subscribe, Topic, SubscriberPid}, _From, State = #state{reverse_tid = ReverseTid}) ->
    Components = of_components(Topic),
    case is_valid_components(Components) of
        true ->
            case has_subscription(ReverseTid, Topic, SubscriberPid) of
                true ->
                    {reply, ok, State};
                false ->
                    {MonitorRef, State1} = ensure_pid_monitor(SubscriberPid, State),
                    Sub = #subscriber{
                        topic = Topic,
                        subscriber_pid = SubscriberPid,
                        components = Components,
                        monitor_ref = MonitorRef,
                        order = order_num(Components)
                    },
                    State2 = insert_subscription(Topic, Components, SubscriberPid, Sub, State1),
                    {reply, ok, State2}
            end;
        false ->
            {reply, {error, <<"invalid topic name">>}, State}
    end;

handle_call({unsubscribe, Topic, SubscriberPid}, _From, State = #state{reverse_tid = ReverseTid}) ->
    Removed = [Reverse || Reverse = {SubscriberPid0, Topic0, _Kind, _NodeId, _Sub} <- ets:lookup(ReverseTid, SubscriberPid),
        SubscriberPid =:= SubscriberPid0, Topic =:= Topic0],
    lists:foreach(fun(Reverse) -> delete_subscription(Reverse, State) end, Removed),
    State1 = release_pid_monitor(SubscriberPid, length(Removed), State),
    {reply, ok, State1}.

%% @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(_Request, State = #state{}) ->
    {noreply, State}.

%% @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({'DOWN', MonitorRef, process, SubscriberPid, Reason}, State = #state{reverse_tid = ReverseTid, pid_tid = PidTid}) ->
    logger:debug("[endpoint_subscription] subscriber: ~p, down with reason: ~p", [SubscriberPid, Reason]),
    Removed = ets:lookup(ReverseTid, SubscriberPid),
    lists:foreach(fun(Reverse) -> delete_subscription(Reverse, State) end, Removed),
    case ets:lookup(PidTid, SubscriberPid) of
        [{SubscriberPid, MonitorRef, _Count}] ->
            ets:delete(PidTid, SubscriberPid);
        _ ->
            ok
    end,
    {noreply, State};

handle_info(Info, State = #state{}) ->
    logger:debug("[endpoint_subscription] get unknown info: ~p", [Info]),
    {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
%%%===================================================================

-spec match_route_key(binary()) -> [#subscriber{}].
match_route_key(RouteKey) when is_binary(RouteKey) ->
    ExactSubs = [Sub || {_Topic, Sub} <- ets:lookup(?EXACT_TAB, RouteKey)],
    TrieSubs = match_trie(of_components(RouteKey)),
    Sorted = lists:sort(fun compare_subscriber/2, ExactSubs ++ TrieSubs),
    dedupe_subscribers(Sorted).

-spec maybe_log_unmatched_publish(binary(), binary(), [#subscriber{}]) -> ok.
maybe_log_unmatched_publish(RouteKey, Content, []) ->
    endpoint_log:unmatched_publish(RouteKey, Content);
maybe_log_unmatched_publish(_RouteKey, _Content, _MatchedSubscribers) ->
    ok.

-spec of_components(Topic :: binary()) -> [binary()].
of_components(Topic) when is_binary(Topic) ->
    binary:split(Topic, <<$/>>, [global]).

is_valid_components([]) ->
    true;
is_valid_components([<<$+>>|T]) ->
    length(T) =:= 0;
is_valid_components([<<$*>>|T]) ->
    is_valid_components(T);
is_valid_components([_|T]) ->
    is_valid_components(T).

-spec order_num(Components :: list()) -> integer().
order_num([]) ->
    1;
order_num([<<$*>>|_]) ->
    2;
order_num([<<$+>>|_]) ->
    3;
order_num([_|Tail]) ->
    order_num(Tail).

-spec has_subscription(ets:tid(), binary(), pid()) -> boolean().
has_subscription(ReverseTid, Topic, SubscriberPid) ->
    Subscriptions = ets:lookup(ReverseTid, SubscriberPid),
    lists:any(fun({SubscriberPid0, Topic0, _Kind, _NodeId, _Sub}) ->
        Topic =:= Topic0 andalso SubscriberPid =:= SubscriberPid0
    end, Subscriptions).

-spec compare_subscriber(#subscriber{}, #subscriber{}) -> boolean().
compare_subscriber(#subscriber{order = Order0, topic = Topic0}, #subscriber{order = Order1, topic = Topic1}) ->
    case Order0 =:= Order1 of
        true ->
            Topic0 =< Topic1;
        false ->
            Order0 < Order1
    end.

-spec dedupe_subscribers([#subscriber{}]) -> [#subscriber{}].
dedupe_subscribers(Subscribers) ->
    {_, Result} = lists:foldl(fun(S = #subscriber{subscriber_pid = SubscriberPid}, {Seen, Acc}) ->
        case sets:is_element(SubscriberPid, Seen) of
            true ->
                {Seen, Acc};
            false ->
                {sets:add_element(SubscriberPid, Seen), [S | Acc]}
        end
    end, {sets:new(), []}, Subscribers),
    lists:reverse(Result).

-spec insert_subscription(binary(), [binary()], pid(), #subscriber{}, #state{}) -> #state{}.
insert_subscription(Topic, Components, SubscriberPid, Sub, State = #state{exact_tid = ExactTid, trie_sub_tid = TrieSubTid, reverse_tid = ReverseTid}) ->
    case has_wildcard(Components) of
        false ->
            true = ets:insert(ExactTid, {Topic, Sub}),
            true = ets:insert(ReverseTid, {SubscriberPid, Topic, exact, undefined, Sub}),
            State;
        true ->
            {Kind, NodeId, State1} = insert_trie_subscription(Components, State),
            true = ets:insert(TrieSubTid, {{NodeId, subscription_match_type(Kind)}, Sub}),
            true = ets:insert(ReverseTid, {SubscriberPid, Topic, Kind, NodeId, Sub}),
            State1
    end.

-spec insert_trie_subscription([binary()], #state{}) -> {trie_exact | trie_plus, non_neg_integer(), #state{}}.
insert_trie_subscription(Components, State) ->
    case Components of
        [] ->
            {trie_exact, ?ROOT_NODE, State};
        _ ->
            case lists:last(Components) of
                <<"+">> ->
                    Prefix = lists:sublist(Components, length(Components) - 1),
                    {NodeId, State1} = ensure_path(Prefix, State),
                    {trie_plus, NodeId, State1};
                _ ->
                    {NodeId, State1} = ensure_path(Components, State),
                    {trie_exact, NodeId, State1}
            end
    end.

-spec subscription_match_type(trie_exact | trie_plus) -> exact | plus.
subscription_match_type(trie_exact) ->
    exact;
subscription_match_type(trie_plus) ->
    plus.

-spec ensure_path([binary()], #state{}) -> {non_neg_integer(), #state{}}.
ensure_path(Components, State) ->
    ensure_path(?ROOT_NODE, Components, State).

-spec ensure_path(non_neg_integer(), [binary()], #state{}) -> {non_neg_integer(), #state{}}.
ensure_path(NodeId, [], State) ->
    {NodeId, State};
ensure_path(NodeId, [Segment | Rest], State = #state{edge_tid = EdgeTid, next_node_id = NextNodeId}) ->
    EdgeKey = {NodeId, Segment},
    case ets:lookup(EdgeTid, EdgeKey) of
        [{EdgeKey, ChildNodeId}] ->
            ensure_path(ChildNodeId, Rest, State);
        [] ->
            true = ets:insert(EdgeTid, {EdgeKey, NextNodeId}),
            ensure_path(NextNodeId, Rest, State#state{next_node_id = NextNodeId + 1})
    end.

-spec delete_subscription(tuple(), #state{}) -> ok.
delete_subscription(Reverse = {_SubscriberPid, Topic, exact, undefined, Sub}, #state{exact_tid = ExactTid, reverse_tid = ReverseTid}) ->
    true = ets:delete_object(ExactTid, {Topic, Sub}),
    true = ets:delete_object(ReverseTid, Reverse),
    ok;
delete_subscription(Reverse = {_SubscriberPid, _Topic, Kind, NodeId, Sub}, #state{trie_sub_tid = TrieSubTid, reverse_tid = ReverseTid}) ->
    true = ets:delete_object(TrieSubTid, {{NodeId, subscription_match_type(Kind)}, Sub}),
    true = ets:delete_object(ReverseTid, Reverse),
    ok.

-spec ensure_pid_monitor(pid(), #state{}) -> {reference(), #state{}}.
ensure_pid_monitor(SubscriberPid, State = #state{pid_tid = PidTid}) ->
    case ets:lookup(PidTid, SubscriberPid) of
        [{SubscriberPid, MonitorRef, Count}] ->
            true = ets:insert(PidTid, {SubscriberPid, MonitorRef, Count + 1}),
            {MonitorRef, State};
        [] ->
            MonitorRef = erlang:monitor(process, SubscriberPid),
            true = ets:insert(PidTid, {SubscriberPid, MonitorRef, 1}),
            {MonitorRef, State}
    end.

-spec release_pid_monitor(pid(), non_neg_integer(), #state{}) -> #state{}.
release_pid_monitor(_SubscriberPid, 0, State) ->
    State;
release_pid_monitor(SubscriberPid, RemovedCount, State = #state{pid_tid = PidTid}) ->
    case ets:lookup(PidTid, SubscriberPid) of
        [{SubscriberPid, MonitorRef, Count}] when Count =< RemovedCount ->
            erlang:demonitor(MonitorRef, [flush]),
            ets:delete(PidTid, SubscriberPid),
            State;
        [{SubscriberPid, MonitorRef, Count}] ->
            true = ets:insert(PidTid, {SubscriberPid, MonitorRef, Count - RemovedCount}),
            State;
        [] ->
            State
    end.

-spec match_trie([binary()]) -> [#subscriber{}].
match_trie(Components) ->
    match_trie([?ROOT_NODE], Components, []).

-spec match_trie([non_neg_integer()], [binary()], [#subscriber{}]) -> [#subscriber{}].
match_trie(NodeIds, [], Acc) ->
    collect_trie_subscribers(NodeIds, exact) ++ Acc;
match_trie(NodeIds, [Segment | Rest], Acc) ->
    PlusSubs = collect_trie_subscribers(NodeIds, plus),
    NextNodeIds = next_trie_nodes(NodeIds, Segment),
    match_trie(NextNodeIds, Rest, PlusSubs ++ Acc).

-spec collect_trie_subscribers([non_neg_integer()], exact | plus) -> [#subscriber{}].
collect_trie_subscribers(NodeIds, MatchType) ->
    lists:flatmap(fun(NodeId) ->
        [Sub || {{NodeId0, MatchType0}, Sub} <- ets:lookup(?TRIE_SUB_TAB, {NodeId, MatchType}),
            NodeId0 =:= NodeId, MatchType0 =:= MatchType]
    end, NodeIds).

-spec next_trie_nodes([non_neg_integer()], binary()) -> [non_neg_integer()].
next_trie_nodes(NodeIds, Segment) ->
    lists:usort(lists:flatmap(fun(NodeId) ->
        lookup_child(NodeId, Segment) ++ lookup_child(NodeId, <<"*">>)
    end, NodeIds)).

-spec lookup_child(non_neg_integer(), binary()) -> [non_neg_integer()].
lookup_child(NodeId, Segment) ->
    EdgeKey = {NodeId, Segment},
    case ets:lookup(?TRIE_EDGE_TAB, EdgeKey) of
        [{EdgeKey, ChildNodeId}] ->
            [ChildNodeId];
        [] ->
            []
    end.

-spec has_wildcard([binary()]) -> boolean().
has_wildcard(Components) ->
    lists:any(fun(Component) -> Component =:= <<"*">> orelse Component =:= <<"+">> end, Components).

-spec exact_subscribers(ets:tid()) -> [#subscriber{}].
exact_subscribers(ExactTid) ->
    [Sub || {_Topic, Sub} <- ets:tab2list(ExactTid)].

-spec trie_subscribers(ets:tid()) -> [#subscriber{}].
trie_subscribers(TrieSubTid) ->
    [Sub || {_Key, Sub} <- ets:tab2list(TrieSubTid)].
