%%%-------------------------------------------------------------------
%%% @doc
%%% DHT operations for pub/sub - handles subscription advertisement and discovery.
%%%
%%% Responsibilities:
%%% - Advertise subscriptions in DHT with TTL
%%% - Schedule re-advertisement timers
%%% - Discover remote subscribers via DHT queries
%%% - Route messages to remote subscribers
%%% - Track pending DHT queries
%%%
%%% Extracted from macula_pubsub_handler.erl (Phase 3)
%%% @end
%%%-------------------------------------------------------------------
-module(macula_pubsub_dht).

-include_lib("kernel/include/logger.hrl").

%% API
-export([advertise_subscription/5, cancel_advertisement/2,
         discover_subscribers/6, handle_discovery_response/3,
         route_to_subscribers/5]).

-type topic() :: binary().
-type node_id() :: binary().
-type url() :: binary().
-type connection_manager_pid() :: pid().
-type subscription_ref() :: reference().
-type payload() :: binary().
-type qos() :: 0 | 1.

-type advertised_subscriptions() :: #{topic() => #{
    sub_ref := reference(),
    ttl := pos_integer(),
    timer_ref := reference()
}}.

-type pending_queries() :: #{binary() => {topic(), payload(), qos(), map()}}.

-export_type([advertised_subscriptions/0, pending_queries/0]).

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

%% @doc Advertise a subscription in the DHT.
%% Sends STORE message to DHT and schedules re-advertisement.
%% Returns {ok, SubInfo}.
-spec advertise_subscription(topic(), subscription_ref(), node_id(), url(), connection_manager_pid()) ->
    {ok, #{sub_ref := reference(), ttl := pos_integer(), timer_ref := reference()}} | {error, term()}.
advertise_subscription(Topic, SubRef, NodeId, Url, _ConnMgrPid) ->
    %% Default TTL for subscriptions is 300 seconds (5 minutes)
    TTL = 300,

    %% Create DHT key by hashing the topic
    TopicKey = crypto:hash(sha256, Topic),

    %% Create subscriber value with endpoint info for routing messages back
    SubscriberValue = #{
        node_id => NodeId,
        endpoint => Url,
        ttl => TTL
    },

    %% Store subscription in DHT with propagation to k closest nodes (v0.8.0+)
    ?LOG_INFO("[~s] Advertising subscription to topic ~s in DHT", [NodeId, Topic]),
    try
        case whereis(macula_routing_server) of
            undefined ->
                ?LOG_WARNING("[~s] Routing server not running, cannot advertise subscription", [NodeId]);
            RoutingServerPid ->
                case macula_routing_server:store(RoutingServerPid, TopicKey, SubscriberValue) of
                    ok ->
                        ?LOG_DEBUG("[~s] Successfully stored subscription for ~s in DHT", [NodeId, Topic]);
                    {error, StoreError} ->
                        ?LOG_WARNING("[~s] Failed to store subscription ~s: ~p",
                                    [NodeId, Topic, StoreError])
                end
        end
    catch
        _:DhtError ->
            ?LOG_WARNING("[~s] Failed to advertise subscription ~s in DHT: ~p (continuing)",
                        [NodeId, Topic, DhtError])
    end,

    %% Schedule re-advertisement (TTL - 60 seconds, minimum 10 seconds)
    ResubInterval = max(10, TTL - 60) * 1000,  % milliseconds
    TimerRef = erlang:send_after(ResubInterval, self(), {resubscribe, Topic}),
    ?LOG_DEBUG("[~s] Scheduled re-subscription for ~s in ~p seconds",
              [NodeId, Topic, ResubInterval div 1000]),

    %% Return subscription info
    SubInfo = #{
        sub_ref => SubRef,
        ttl => TTL,
        timer_ref => TimerRef
    },
    {ok, SubInfo}.

%% @doc Cancel advertisement for a topic.
%% Cancels the re-advertisement timer.
%% Returns updated advertised_subscriptions map.
-spec cancel_advertisement(topic(), advertised_subscriptions()) -> advertised_subscriptions().
cancel_advertisement(Topic, AdvertisedSubscriptions) ->
    case maps:get(Topic, AdvertisedSubscriptions, undefined) of
        undefined ->
            %% Not advertised, return unchanged
            AdvertisedSubscriptions;
        SubInfo ->
            %% Cancel timer and remove from map
            TimerRef = maps:get(timer_ref, SubInfo),
            erlang:cancel_timer(TimerRef),
            maps:remove(Topic, AdvertisedSubscriptions)
    end.

%% @doc Discover remote subscribers for a topic.
%% Checks cache first, queries DHT on cache miss.
%% Returns {cached, Subscribers, Registry} | {query_sent, Pending, MsgId, Registry}.
-spec discover_subscribers(topic(), payload(), qos(), connection_manager_pid(), term(), non_neg_integer()) ->
    {cached, list(), term()} | {query_sent, pending_queries(), binary(), term()}.
discover_subscribers(Topic, Payload, Qos, ConnMgrPid, ServiceRegistry, MsgIdCounter) ->
    %% Check subscriber cache first
    case macula_service_registry:discover_subscribers(ServiceRegistry, Topic) of
        {ok, Subscribers, UpdatedRegistry} ->
            %% Cache hit - return cached subscribers
            ?LOG_DEBUG("Cache hit for subscribers to topic: ~s (~p subscribers)",
                      [Topic, length(Subscribers)]),
            {cached, Subscribers, UpdatedRegistry};

        {cache_miss, UpdatedRegistry} ->
            %% Cache miss - query DHT asynchronously
            {MsgId, _NewCounter} = next_message_id(MsgIdCounter),
            query_dht_async(Topic, Payload, Qos, MsgId, ConnMgrPid),

            %% Track pending query
            Pending = #{MsgId => {Topic, Payload, Qos, #{}}},
            {query_sent, Pending, MsgId, UpdatedRegistry}
    end.

%% @doc Handle DHT discovery response.
%% Routes messages to discovered subscribers.
%% Returns updated pending queries map.
-spec handle_discovery_response(binary(), list(), pending_queries()) ->
    {ok, pending_queries()} | {not_found, pending_queries()}.
handle_discovery_response(MsgId, _Subscribers, PendingQueries) ->
    case maps:get(MsgId, PendingQueries, undefined) of
        undefined ->
            %% Query not found (already handled or unknown)
            {not_found, PendingQueries};
        {_Topic, _Payload, _Qos, _Opts} ->
            %% Query found - remove from pending
            %% Note: Actual routing happens in the caller
            {ok, maps:remove(MsgId, PendingQueries)}
    end.

%% @doc Route message to remote subscribers via direct P2P connections (v0.8.0+).
%% Wraps publish in pubsub_route envelope and sends directly to each subscriber.
%% Uses macula_peer_connector for direct QUIC connections to subscriber endpoints.
-spec route_to_subscribers(topic(), payload(), qos(), list(), node_id()) -> ok.
route_to_subscribers(_Topic, _Payload, _Qos, [], _NodeId) ->
    ok;
route_to_subscribers(Topic, Payload, Qos, Subscribers, SourceNodeId) ->
    ?LOG_INFO("[~s] Routing message to ~p remote subscriber(s) for topic: ~s via P2P",
             [SourceNodeId, length(Subscribers), Topic]),

    %% Route to each subscriber via direct P2P connection
    lists:foreach(
        fun(Subscriber) ->
            %% Extract subscriber node_id and endpoint
            NodeId = maps:get(node_id, Subscriber, maps:get(<<"node_id">>, Subscriber, undefined)),
            Endpoint = maps:get(endpoint, Subscriber, maps:get(<<"endpoint">>, Subscriber, undefined)),

            case {NodeId, Endpoint} of
                {undefined, _} ->
                    ?LOG_WARNING("[~s] Subscriber missing node_id, skipping", [SourceNodeId]);
                {_, undefined} ->
                    ?LOG_WARNING("[~s] Subscriber missing endpoint, skipping", [SourceNodeId]);
                {DestNodeId, DestEndpoint} ->
                    %% Build PUBLISH message
                    PublishMsg = #{
                        <<"topic">> => Topic,
                        <<"payload">> => Payload,
                        <<"qos">> => Qos,
                        <<"retain">> => false,
                        <<"message_id">> => crypto:strong_rand_bytes(16)
                    },

                    %% Wrap in pubsub_route envelope (MaxHops = 10)
                    PubSubRouteMsg = macula_pubsub_routing:wrap_publish(
                        SourceNodeId, DestNodeId, PublishMsg, 10
                    ),

                    %% Send directly to subscriber via peer connector
                    case macula_peer_connector:send_message(DestEndpoint, pubsub_route, PubSubRouteMsg) of
                        ok ->
                            ?LOG_DEBUG("[~s] Sent pubsub_route directly to subscriber ~s at ~s for topic ~s",
                                      [SourceNodeId, binary:encode_hex(DestNodeId), DestEndpoint, Topic]);
                        {error, Reason} ->
                            ?LOG_ERROR("[~s] Failed to send pubsub_route to ~s: ~p",
                                      [SourceNodeId, DestEndpoint, Reason])
                    end
            end
        end,
        Subscribers
    ),
    ok.

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

%% @doc Generate next message ID
-spec next_message_id(non_neg_integer()) -> {binary(), non_neg_integer()}.
next_message_id(Counter) ->
    macula_utils:next_message_id(Counter).

%% @doc Query DHT for subscribers synchronously (v0.8.0+).
%% Directly queries local routing server instead of sending message to connection manager.
-spec query_dht_async(topic(), payload(), qos(), binary(), connection_manager_pid()) -> ok.
query_dht_async(Topic, _Payload, _Qos, _MsgId, _ConnMgrPid) ->
    %% Create DHT key from topic
    TopicKey = crypto:hash(sha256, Topic),

    ?LOG_DEBUG("Querying DHT for remote subscribers to topic: ~s", [Topic]),

    %% Query local routing server directly (v0.8.0+)
    %% This is actually synchronous now, but we keep the function name for compatibility
    try
        case whereis(macula_routing_server) of
            undefined ->
                ?LOG_WARNING("Routing server not running, cannot query for subscribers to ~s", [Topic]);
            RoutingServerPid ->
                %% K=20 is standard Kademlia replication factor
                case macula_routing_server:find_value(RoutingServerPid, TopicKey, 20) of
                    {ok, Subscribers} when is_list(Subscribers), length(Subscribers) > 0 ->
                        ?LOG_DEBUG("Found ~p subscriber(s) for topic ~s in DHT",
                                  [length(Subscribers), Topic]);
                    {ok, []} ->
                        ?LOG_DEBUG("No subscribers found for topic ~s in DHT", [Topic]);
                    {error, not_found} ->
                        ?LOG_DEBUG("No subscribers found for topic ~s in DHT", [Topic]);
                    {error, QueryError} ->
                        ?LOG_WARNING("Failed to query DHT for topic ~s: ~p",
                                    [Topic, QueryError])
                end
        end
    catch
        _:Error:Stack ->
            ?LOG_WARNING("Failed to query DHT for topic ~s: ~p~nStack: ~p",
                        [Topic, Error, Stack])
    end,
    ok.