%%%-------------------------------------------------------------------
%%% @doc Macula Ping/Pong - Direct P2P Async RPC Demo
%%%
%%% A peer-to-peer communication module that demonstrates bidirectional
%%% messaging across NAT boundaries using Macula's async RPC with
%%% direct P2P delivery (NATS-style request/reply).
%%%
%%% Each node:
%%% - Registers local "ping.handler" RPC handler
%%% - Periodically sends PING RPCs to random peers (direct P2P)
%%% - Receives PONGs via callback
%%% - Measures Round-Trip Time (RTT) for each exchange
%%% - Tracks statistics per peer and per NAT type
%%%
%%% Architecture:
%%% - Bootstrap: DHT only (no pub/sub routing)
%%% - Communication: Direct P2P via QUIC
%%% - Discovery: DHT for service registration
%%%
%%% Flow:
%%% 1. fc01 registers local "ping.handler" (handles incoming PINGs)
%%% 2. fc01 calls macula_rpc_handler:request/4 to ping rc05
%%% 3. Request goes directly to rc05 via QUIC (NAT-aware)
%%% 4. rc05's "ping.handler" executes, returns PONG
%%% 5. fc01 receives PONG via callback, measures RTT
%%%
%%% @end
%%%-------------------------------------------------------------------
-module(macula_ping_pong).

-behaviour(gen_server).

%% API
-export([start_link/0, start_link/1]).
-export([ping/1, get_stats/0, get_peer_stats/1, reset_stats/0]).

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

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

-define(SERVER, ?MODULE).
-define(DEFAULT_INTERVAL, 5000).      % 5 seconds between PINGs
-define(DEFAULT_TIMEOUT, 3000).       % 3 seconds PING timeout
-define(PING_PROCEDURE, <<"ping.handler">>).

-record(peer_stats, {
    pings_sent = 0 :: non_neg_integer(),
    pongs_received = 0 :: non_neg_integer(),
    timeouts = 0 :: non_neg_integer(),
    min_rtt_ms :: non_neg_integer() | undefined,
    max_rtt_ms :: non_neg_integer() | undefined,
    total_rtt_ms = 0 :: non_neg_integer(),
    last_seen :: integer() | undefined,
    nat_type :: binary() | undefined
}).

-record(pending_ping, {
    target_peer :: binary(),
    sent_at :: integer(),
    from :: term() | undefined  % For manual ping calls
}).

-record(state, {
    node_id :: binary(),
    nat_type :: binary(),
    known_peers = [] :: [binary()],
    pending_pings = #{} :: #{binary() => #pending_ping{}},
    peer_stats = #{} :: #{binary() => #peer_stats{}},
    global_stats :: #peer_stats{},
    interval :: pos_integer(),
    timeout :: pos_integer(),
    timer_ref :: reference() | undefined,
    setup_complete = false :: boolean()
}).

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

%% @doc Start with default settings
-spec start_link() -> {ok, pid()} | {error, term()}.
start_link() ->
    start_link(#{}).

%% @doc Start with options
%% Options:
%%   - interval: milliseconds between PINGs (default: 5000)
%%   - timeout: milliseconds to wait for PONG (default: 3000)
%%   - node_id: custom node identifier (default: NODE_ID env var)
-spec start_link(map()) -> {ok, pid()} | {error, term()}.
start_link(Opts) ->
    gen_server:start_link({local, ?SERVER}, ?MODULE, Opts, []).

%% @doc Send a PING to a specific peer
-spec ping(binary()) -> {ok, integer()} | {error, term()}.
ping(PeerId) ->
    gen_server:call(?SERVER, {ping, PeerId}, 10000).

%% @doc Get all statistics
-spec get_stats() -> map().
get_stats() ->
    gen_server:call(?SERVER, get_stats).

%% @doc Get statistics for a specific peer
-spec get_peer_stats(binary()) -> map() | not_found.
get_peer_stats(PeerId) ->
    gen_server:call(?SERVER, {get_peer_stats, PeerId}).

%% @doc Reset all statistics
-spec reset_stats() -> ok.
reset_stats() ->
    gen_server:cast(?SERVER, reset_stats).

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

init(Opts) ->
    NodeId = get_node_id(Opts),
    NatType = get_nat_type(NodeId),
    Interval = maps:get(interval, Opts, ?DEFAULT_INTERVAL),
    Timeout = maps:get(timeout, Opts, ?DEFAULT_TIMEOUT),

    io:format("[PingPong ~s] Starting up (NAT: ~s)...~n", [NodeId, NatType]),

    %% Schedule setup after gen_server is fully started
    self() ! setup,

    {ok, #state{
        node_id = NodeId,
        nat_type = NatType,
        interval = Interval,
        timeout = Timeout,
        global_stats = #peer_stats{}
    }}.

handle_call({ping, PeerId}, From, State) ->
    %% Manual ping request - async response
    NewState = do_send_ping(PeerId, From, State),
    {noreply, NewState};

handle_call(get_stats, _From, State) ->
    Stats = format_all_stats(State),
    {reply, Stats, State};

handle_call({get_peer_stats, PeerId}, _From, State) ->
    case maps:get(PeerId, State#state.peer_stats, not_found) of
        not_found ->
            {reply, not_found, State};
        PeerStats ->
            {reply, format_peer_stats(PeerId, PeerStats), State}
    end;

handle_call(_Request, _From, State) ->
    {reply, {error, unknown_call}, State}.

handle_cast(reset_stats, State) ->
    io:format("[PingPong ~s] Resetting statistics~n", [State#state.node_id]),
    {noreply, State#state{
        peer_stats = #{},
        global_stats = #peer_stats{}
    }};

handle_cast(_Msg, State) ->
    {noreply, State}.

handle_info(setup, State) ->
    NewState = setup_ping_pong(State),
    {noreply, NewState};

handle_info(ping_tick, State) ->
    NewState = do_ping_tick(State),
    %% Schedule next tick
    TimerRef = erlang:send_after(State#state.interval, self(), ping_tick),
    {noreply, NewState#state{timer_ref = TimerRef}};

%% Handle async RPC reply via process message
handle_info({rpc_reply, RequestId, Result}, State) ->
    NewState = handle_rpc_reply(RequestId, Result, State),
    {noreply, NewState};

handle_info({peer_discovered, PeerId}, State) when PeerId =/= State#state.node_id ->
    %% Add to known peers if not already present
    KnownPeers = State#state.known_peers,
    NewKnownPeers = case lists:member(PeerId, KnownPeers) of
        true -> KnownPeers;
        false ->
            io:format("[PingPong ~s] Discovered peer: ~s~n", [State#state.node_id, PeerId]),
            [PeerId | KnownPeers]
    end,
    {noreply, State#state{known_peers = NewKnownPeers}};

handle_info({peer_discovered, _SelfId}, State) ->
    %% Ignore self-discovery
    {noreply, State};

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

terminate(_Reason, State) ->
    io:format("[PingPong ~s] Shutting down~n", [State#state.node_id]),
    print_final_stats(State),
    ok.

%%%===================================================================
%%% Internal functions - Setup
%%%===================================================================

%% @private Get node identifier
get_node_id(Opts) ->
    case maps:get(node_id, Opts, undefined) of
        undefined ->
            case os:getenv("NODE_ID") of
                false ->
                    {ok, Hostname} = inet:gethostname(),
                    list_to_binary(Hostname);
                NodeId ->
                    list_to_binary(NodeId)
            end;
        NodeId when is_binary(NodeId) ->
            NodeId;
        NodeId when is_list(NodeId) ->
            list_to_binary(NodeId)
    end.

%% @private Determine NAT type from node ID prefix
get_nat_type(NodeId) ->
    case NodeId of
        <<"fc", _/binary>> -> <<"full_cone">>;
        <<"rc", _/binary>> -> <<"restricted">>;
        <<"sy", _/binary>> -> <<"symmetric">>;
        <<"bootstrap", _/binary>> -> <<"public">>;
        _ -> <<"unknown">>
    end.

%% @private Setup RPC handler and start ping loop
setup_ping_pong(State) ->
    NodeId = State#state.node_id,
    NatType = State#state.nat_type,
    io:format("[PingPong ~s] Setting up async RPC handler...~n", [NodeId]),

    %% Wait for mesh to stabilize
    timer:sleep(3000),

    %% Register local PING handler
    case register_ping_handler(NodeId, NatType) of
        ok ->
            io:format("[PingPong ~s] Registered ~s handler~n", [NodeId, ?PING_PROCEDURE]);
        {error, Reason} ->
            io:format("[PingPong ~s] Failed to register handler: ~p~n", [NodeId, Reason])
    end,

    %% Discover initial peers
    KnownPeers = discover_peers(State),
    io:format("[PingPong ~s] Discovered ~p initial peers~n", [NodeId, length(KnownPeers)]),

    %% Start ping loop
    io:format("[PingPong ~s] Starting PING loop (interval: ~pms, timeout: ~pms)~n",
              [NodeId, State#state.interval, State#state.timeout]),
    TimerRef = erlang:send_after(State#state.interval, self(), ping_tick),

    State#state{
        timer_ref = TimerRef,
        known_peers = KnownPeers,
        setup_complete = true
    }.

%% @private Register local handler for incoming PINGs
register_ping_handler(NodeId, NatType) ->
    %% Handler receives PING args and returns PONG
    Handler = fun(Args) ->
        FromNode = maps:get(<<"from">>, Args, <<"unknown">>),
        SentAt = maps:get(<<"sent_at">>, Args, 0),

        ?LOG_DEBUG("PING handler invoked from ~s", [FromNode]),

        %% Return PONG response
        {ok, #{
            <<"type">> => <<"pong">>,
            <<"from">> => NodeId,
            <<"nat_type">> => NatType,
            <<"sent_at">> => SentAt,
            <<"received_at">> => erlang:system_time(millisecond),
            <<"responded_at">> => erlang:system_time(millisecond)
        }}
    end,

    %% Register with local RPC handler
    case get_rpc_handler() of
        {ok, RpcPid} ->
            macula_rpc_handler:register_local_procedure(RpcPid, ?PING_PROCEDURE, Handler),
            ok;
        {error, Reason} ->
            {error, Reason}
    end.

%% @private Discover available peers
discover_peers(State) ->
    NodeId = State#state.node_id,

    %% Generate list of expected peer IDs based on 50-peer setup
    FullCone = [iolist_to_binary([<<"fc">>, pad_num(N)]) || N <- lists:seq(1, 17)],
    Restricted = [iolist_to_binary([<<"rc">>, pad_num(N)]) || N <- lists:seq(1, 17)],
    Symmetric = [iolist_to_binary([<<"sy">>, pad_num(N)]) || N <- lists:seq(1, 16)],

    AllPeers = FullCone ++ Restricted ++ Symmetric,

    %% Remove self
    lists:filter(fun(P) -> P =/= NodeId end, AllPeers).

%% @private Pad number to 2 digits
pad_num(N) when N < 10 ->
    iolist_to_binary([<<"0">>, integer_to_binary(N)]);
pad_num(N) ->
    integer_to_binary(N).

%%%===================================================================
%%% Internal functions - PING/PONG
%%%===================================================================

%% @private Handle periodic ping tick
do_ping_tick(State) ->
    case select_random_peer(State) of
        {ok, TargetPeer} ->
            do_send_ping(TargetPeer, undefined, State);
        {error, no_peers} ->
            io:format("[PingPong ~s] No peers available to ping~n", [State#state.node_id]),
            State
    end.

%% @private Select a random peer to ping
select_random_peer(#state{known_peers = []}) ->
    {error, no_peers};
select_random_peer(#state{known_peers = Peers}) ->
    Index = rand:uniform(length(Peers)),
    {ok, lists:nth(Index, Peers)}.

%% @private Send a PING to target peer using async RPC
do_send_ping(TargetPeer, From, State) ->
    NodeId = State#state.node_id,
    NatType = State#state.nat_type,
    SentAt = erlang:system_time(millisecond),

    io:format("[PingPong ~s] PING -> ~s (async RPC)~n", [NodeId, TargetPeer]),

    %% Build PING args
    PingArgs = #{
        <<"type">> => <<"ping">>,
        <<"from">> => NodeId,
        <<"nat_type">> => NatType,
        <<"sent_at">> => SentAt
    },

    %% Send async RPC request
    case get_rpc_handler() of
        {ok, RpcPid} ->
            Opts = #{timeout => State#state.timeout},
            case macula_rpc_handler:request(RpcPid, ?PING_PROCEDURE, PingArgs, Opts) of
                {ok, RequestId} ->
                    %% Track pending PING
                    PendingInfo = #pending_ping{
                        target_peer = TargetPeer,
                        sent_at = SentAt,
                        from = From
                    },
                    NewPending = maps:put(RequestId, PendingInfo, State#state.pending_pings),

                    %% Update peer stats (ping sent)
                    NewPeerStats = update_ping_sent(TargetPeer, State#state.peer_stats),
                    NewGlobal = increment_sent(State#state.global_stats),

                    State#state{
                        pending_pings = NewPending,
                        peer_stats = NewPeerStats,
                        global_stats = NewGlobal
                    };
                {error, Reason} ->
                    io:format("[PingPong ~s] RPC request failed: ~p~n", [NodeId, Reason]),
                    maybe_reply(From, {error, Reason}),
                    State
            end;
        {error, Reason} ->
            io:format("[PingPong ~s] No RPC handler: ~p~n", [NodeId, Reason]),
            maybe_reply(From, {error, Reason}),
            State
    end.

%% @private Handle async RPC reply
handle_rpc_reply(RequestId, Result, State) ->
    case maps:take(RequestId, State#state.pending_pings) of
        {#pending_ping{target_peer = TargetPeer, sent_at = SentAt, from = From}, NewPending} ->
            case Result of
                {ok, PongMsg} ->
                    %% Calculate RTT
                    ReceivedAt = erlang:system_time(millisecond),
                    RTT = ReceivedAt - SentAt,

                    %% Extract NAT type from response
                    PeerNat = maps:get(<<"nat_type">>, PongMsg, <<"unknown">>),

                    io:format("[PingPong ~s] PONG <- ~s (~s): ~pms RTT~n",
                              [State#state.node_id, TargetPeer, PeerNat, RTT]),

                    %% Update statistics
                    NewPeerStats = update_pong_received(TargetPeer, RTT, PeerNat, State#state.peer_stats),
                    NewGlobal = update_rtt(RTT, State#state.global_stats),

                    %% Reply to caller if manual ping
                    maybe_reply(From, {ok, RTT}),

                    State#state{
                        pending_pings = NewPending,
                        peer_stats = NewPeerStats,
                        global_stats = NewGlobal
                    };

                {error, timeout} ->
                    io:format("[PingPong ~s] PING timeout -> ~s~n",
                              [State#state.node_id, TargetPeer]),

                    %% Update statistics
                    NewPeerStats = update_timeout(TargetPeer, State#state.peer_stats),
                    NewGlobal = increment_timeout(State#state.global_stats),

                    maybe_reply(From, {error, timeout}),

                    State#state{
                        pending_pings = NewPending,
                        peer_stats = NewPeerStats,
                        global_stats = NewGlobal
                    };

                {error, Reason} ->
                    io:format("[PingPong ~s] PING error -> ~s: ~p~n",
                              [State#state.node_id, TargetPeer, Reason]),
                    maybe_reply(From, {error, Reason}),
                    State#state{pending_pings = NewPending}
            end;
        error ->
            %% Unknown request ID (already handled or never sent)
            State
    end.

%% @private Maybe reply to caller
maybe_reply(undefined, _Result) ->
    ok;
maybe_reply(From, Result) ->
    gen_server:reply(From, Result).

%%%===================================================================
%%% Statistics functions
%%%===================================================================

%% @private Update stats when PING is sent
update_ping_sent(PeerId, PeerStats) ->
    Stats = maps:get(PeerId, PeerStats, #peer_stats{}),
    NewStats = Stats#peer_stats{pings_sent = Stats#peer_stats.pings_sent + 1},
    maps:put(PeerId, NewStats, PeerStats).

%% @private Update stats when PONG is received
update_pong_received(PeerId, RTT, NatType, PeerStats) ->
    Stats = maps:get(PeerId, PeerStats, #peer_stats{}),
    NewStats = Stats#peer_stats{
        pongs_received = Stats#peer_stats.pongs_received + 1,
        min_rtt_ms = min_val(Stats#peer_stats.min_rtt_ms, RTT),
        max_rtt_ms = max_val(Stats#peer_stats.max_rtt_ms, RTT),
        total_rtt_ms = Stats#peer_stats.total_rtt_ms + RTT,
        last_seen = erlang:system_time(millisecond),
        nat_type = NatType
    },
    maps:put(PeerId, NewStats, PeerStats).

%% @private Update stats when timeout occurs
update_timeout(PeerId, PeerStats) ->
    Stats = maps:get(PeerId, PeerStats, #peer_stats{}),
    NewStats = Stats#peer_stats{timeouts = Stats#peer_stats.timeouts + 1},
    maps:put(PeerId, NewStats, PeerStats).

%% @private Update global RTT stats
update_rtt(RTT, Stats) ->
    Stats#peer_stats{
        pongs_received = Stats#peer_stats.pongs_received + 1,
        min_rtt_ms = min_val(Stats#peer_stats.min_rtt_ms, RTT),
        max_rtt_ms = max_val(Stats#peer_stats.max_rtt_ms, RTT),
        total_rtt_ms = Stats#peer_stats.total_rtt_ms + RTT
    }.

%% @private Increment sent counter
increment_sent(Stats) ->
    Stats#peer_stats{pings_sent = Stats#peer_stats.pings_sent + 1}.

%% @private Increment timeout counter
increment_timeout(Stats) ->
    Stats#peer_stats{timeouts = Stats#peer_stats.timeouts + 1}.

%% @private Min value handling undefined
min_val(undefined, V) -> V;
min_val(V1, V2) -> min(V1, V2).

%% @private Max value handling undefined
max_val(undefined, V) -> V;
max_val(V1, V2) -> max(V1, V2).

%%%===================================================================
%%% Formatting functions
%%%===================================================================

%% @private Format all statistics
format_all_stats(State) ->
    GlobalStats = State#state.global_stats,
    PeerStatsMap = State#state.peer_stats,

    %% Group by NAT type
    ByNat = group_by_nat(maps:to_list(PeerStatsMap)),

    #{
        node_id => State#state.node_id,
        nat_type => State#state.nat_type,
        global => #{
            pings_sent => GlobalStats#peer_stats.pings_sent,
            pongs_received => GlobalStats#peer_stats.pongs_received,
            timeouts => GlobalStats#peer_stats.timeouts,
            success_rate => calc_success_rate(GlobalStats),
            min_rtt_ms => GlobalStats#peer_stats.min_rtt_ms,
            max_rtt_ms => GlobalStats#peer_stats.max_rtt_ms,
            avg_rtt_ms => calc_avg_rtt(GlobalStats)
        },
        by_nat_type => ByNat,
        peers_contacted => maps:size(PeerStatsMap),
        known_peers => length(State#state.known_peers),
        pending_pings => maps:size(State#state.pending_pings)
    }.

%% @private Format peer statistics
format_peer_stats(PeerId, Stats) ->
    #{
        peer_id => PeerId,
        nat_type => Stats#peer_stats.nat_type,
        pings_sent => Stats#peer_stats.pings_sent,
        pongs_received => Stats#peer_stats.pongs_received,
        timeouts => Stats#peer_stats.timeouts,
        success_rate => calc_success_rate(Stats),
        min_rtt_ms => Stats#peer_stats.min_rtt_ms,
        max_rtt_ms => Stats#peer_stats.max_rtt_ms,
        avg_rtt_ms => calc_avg_rtt(Stats),
        last_seen => Stats#peer_stats.last_seen
    }.

%% @private Group peer stats by NAT type
group_by_nat(PeerStatsList) ->
    lists:foldl(fun({_PeerId, Stats}, Acc) ->
        NatType = case Stats#peer_stats.nat_type of
            undefined -> <<"unknown">>;
            Type -> Type
        end,
        Existing = maps:get(NatType, Acc, #peer_stats{}),
        Merged = merge_stats(Existing, Stats),
        maps:put(NatType, format_nat_group(Merged), Acc)
    end, #{}, PeerStatsList).

%% @private Merge two stats records
merge_stats(A, B) ->
    #peer_stats{
        pings_sent = A#peer_stats.pings_sent + B#peer_stats.pings_sent,
        pongs_received = A#peer_stats.pongs_received + B#peer_stats.pongs_received,
        timeouts = A#peer_stats.timeouts + B#peer_stats.timeouts,
        min_rtt_ms = min_val(A#peer_stats.min_rtt_ms, B#peer_stats.min_rtt_ms),
        max_rtt_ms = max_val(A#peer_stats.max_rtt_ms, B#peer_stats.max_rtt_ms),
        total_rtt_ms = A#peer_stats.total_rtt_ms + B#peer_stats.total_rtt_ms
    }.

%% @private Format NAT group stats
format_nat_group(Stats) ->
    #{
        pings_sent => Stats#peer_stats.pings_sent,
        pongs_received => Stats#peer_stats.pongs_received,
        timeouts => Stats#peer_stats.timeouts,
        success_rate => calc_success_rate(Stats),
        min_rtt_ms => Stats#peer_stats.min_rtt_ms,
        max_rtt_ms => Stats#peer_stats.max_rtt_ms,
        avg_rtt_ms => calc_avg_rtt(Stats)
    }.

%% @private Calculate success rate
calc_success_rate(#peer_stats{pings_sent = 0}) ->
    0.0;
calc_success_rate(#peer_stats{pings_sent = Sent, pongs_received = Received}) ->
    (Received / Sent) * 100.

%% @private Calculate average RTT
calc_avg_rtt(#peer_stats{pongs_received = 0}) ->
    undefined;
calc_avg_rtt(#peer_stats{pongs_received = Count, total_rtt_ms = Total}) ->
    Total / Count.

%% @private Print final statistics on shutdown
print_final_stats(State) ->
    Stats = format_all_stats(State),
    Global = maps:get(global, Stats),

    io:format("~n=== PING/PONG Statistics for ~s (~s) ===~n",
              [maps:get(node_id, Stats), maps:get(nat_type, Stats)]),
    io:format("  PINGs sent:     ~p~n", [maps:get(pings_sent, Global)]),
    io:format("  PONGs received: ~p~n", [maps:get(pongs_received, Global)]),
    io:format("  Timeouts:       ~p~n", [maps:get(timeouts, Global)]),
    io:format("  Success rate:   ~.1f%~n", [maps:get(success_rate, Global)]),
    io:format("  Min RTT:        ~p ms~n", [maps:get(min_rtt_ms, Global)]),
    io:format("  Max RTT:        ~p ms~n", [maps:get(max_rtt_ms, Global)]),
    io:format("  Avg RTT:        ~p ms~n", [maps:get(avg_rtt_ms, Global)]),
    io:format("  Peers contacted: ~p~n", [maps:get(peers_contacted, Stats)]),

    io:format("~n  By NAT Type:~n"),
    maps:foreach(fun(NatType, NatStats) ->
        io:format("    ~s: sent=~p, received=~p, rate=~.1f%, avg_rtt=~p ms~n",
                  [NatType,
                   maps:get(pings_sent, NatStats),
                   maps:get(pongs_received, NatStats),
                   maps:get(success_rate, NatStats),
                   maps:get(avg_rtt_ms, NatStats)])
    end, maps:get(by_nat_type, Stats)),
    io:format("~n").

%%%===================================================================
%%% Peer handler helpers
%%%===================================================================

%% @private Get the RPC handler from the first available peer system
get_rpc_handler() ->
    case whereis(macula_peers_sup) of
        undefined ->
            {error, no_peers_sup};
        _Pid ->
            case macula_peers_sup:list_peers() of
                [] ->
                    {error, no_peers};
                [PeerSystemPid | _] ->
                    get_rpc_from_supervisor(PeerSystemPid)
            end
    end.

%% @private Extract RPC handler PID from peer system supervisor
get_rpc_from_supervisor(SupPid) ->
    case supervisor:which_children(SupPid) of
        Children when is_list(Children) ->
            case find_child_pid(Children, rpc_handler) of
                undefined -> {error, no_rpc_handler};
                RpcPid -> {ok, RpcPid}
            end;
        _ ->
            {error, no_children}
    end.

%% @private Find a child PID by child ID
find_child_pid(Children, ChildId) ->
    case lists:keyfind(ChildId, 1, Children) of
        {ChildId, Pid, _Type, _Modules} when is_pid(Pid) -> Pid;
        _ -> undefined
    end.