%%%-------------------------------------------------------------------
%%% @doc
%%% Connection Upgrade Manager.
%%%
%%% Handles upgrading connections from relay (indirect) to direct
%%% when hole punching succeeds. This is a key optimization for
%%% reducing latency and bootstrap load.
%%%
%%% Upgrade Flow:
%%% 1. Connection established via relay (fallback)
%%% 2. Hole punch attempt happens in background
%%% 3. On success, upgrade_to_direct/2 is called
%%% 4. Messages are seamlessly transitioned to direct connection
%%% 5. Relay connection is gracefully closed
%%%
%%% Key Features:
%%% - Message ordering preserved during upgrade
%%% - No message loss during transition
%%% - Automatic fallback if direct connection fails
%%% - Metrics tracking for upgrade success/failure
%%%
%%% @end
%%%-------------------------------------------------------------------
-module(macula_connection_upgrade).
-behaviour(gen_server).

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

%% API
-export([
    start_link/0,
    upgrade_to_direct/3,
    register_relay/3,
    unregister_relay/1,
    get_upgrade_stats/0
]).

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

-define(SERVER, ?MODULE).

%% Grace period to drain in-flight messages before closing relay
-define(DRAIN_TIMEOUT_MS, 1000).

%%%===================================================================
%%% Types
%%%===================================================================

-record(relay_info, {
    peer_id :: binary(),
    relay_conn :: term(),  % Connection handle to relay
    endpoint :: {inet:ip_address(), inet:port_number()},
    created_at :: integer()
}).

-record(state, {
    %% Map from peer_id -> relay_info
    relays = #{} :: #{binary() => #relay_info{}},
    %% Upgrade statistics
    stats = #{
        upgrades_attempted => 0,
        upgrades_succeeded => 0,
        upgrades_failed => 0
    } :: map()
}).

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

%% @doc Start the connection upgrade manager.
-spec start_link() -> {ok, pid()} | {error, term()}.
start_link() ->
    gen_server:start_link({local, ?SERVER}, ?MODULE, [], []).

%% @doc Upgrade a relay connection to direct.
%% PeerId - The remote peer's ID
%% RelayConn - Current relay connection handle
%% DirectConn - New direct connection handle from hole punch
-spec upgrade_to_direct(binary(), term(), term()) ->
    ok | {error, no_relay | upgrade_failed}.
upgrade_to_direct(PeerId, RelayConn, DirectConn) ->
    gen_server:call(?SERVER, {upgrade, PeerId, RelayConn, DirectConn}, 10000).

%% @doc Register a relay connection for potential future upgrade.
%% Called when connection is established via relay fallback.
-spec register_relay(binary(), term(), {inet:ip_address(), inet:port_number()}) -> ok.
register_relay(PeerId, RelayConn, Endpoint) ->
    gen_server:cast(?SERVER, {register_relay, PeerId, RelayConn, Endpoint}).

%% @doc Unregister a relay connection (closed or upgraded).
-spec unregister_relay(binary()) -> ok.
unregister_relay(PeerId) ->
    gen_server:cast(?SERVER, {unregister_relay, PeerId}).

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

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

init([]) ->
    ?LOG_INFO("Connection upgrade manager started"),
    {ok, #state{}}.

handle_call({upgrade, PeerId, RelayConn, DirectConn}, _From, State) ->
    NewStats = increment_stat(upgrades_attempted, State#state.stats),
    RelayInfo = maps:get(PeerId, State#state.relays, undefined),
    attempt_upgrade(PeerId, RelayConn, DirectConn, RelayInfo, NewStats, State);

handle_call(get_stats, _From, State) ->
    BaseStats = State#state.stats,
    Stats = BaseStats#{active_relays => maps:size(State#state.relays)},
    {reply, Stats, State};

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

handle_cast({register_relay, PeerId, RelayConn, Endpoint}, State) ->
    RelayInfo = #relay_info{
        peer_id = PeerId,
        relay_conn = RelayConn,
        endpoint = Endpoint,
        created_at = erlang:system_time(millisecond)
    },
    ?LOG_DEBUG("Registered relay for peer ~p", [PeerId]),
    NewRelays = maps:put(PeerId, RelayInfo, State#state.relays),
    {noreply, State#state{relays = NewRelays}};

handle_cast({unregister_relay, PeerId}, State) ->
    ?LOG_DEBUG("Unregistered relay for peer ~p", [PeerId]),
    NewRelays = maps:remove(PeerId, State#state.relays),
    {noreply, State#state{relays = NewRelays}};

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

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

terminate(_Reason, _State) ->
    ok.

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

%% @private
%% @doc Attempt upgrade for a peer, dispatching on relay lookup result.
-spec attempt_upgrade(binary(), term(), term(), #relay_info{} | undefined, map(), #state{}) ->
    {reply, term(), #state{}}.
attempt_upgrade(PeerId, _RelayConn, _DirectConn, undefined, Stats, State) ->
    ?LOG_WARNING("Upgrade requested for unknown relay: ~p", [PeerId]),
    {reply, {error, no_relay}, State#state{stats = Stats}};
attempt_upgrade(PeerId, RelayConn, DirectConn, #relay_info{relay_conn = RegisteredRelay} = RelayInfo, Stats, State) ->
    case verify_relay(RelayConn, RegisteredRelay) of
        true ->
            apply_upgrade(PeerId, RelayInfo, DirectConn, Stats, State);
        false ->
            ?LOG_WARNING("Relay connection mismatch for peer ~p", [PeerId]),
            {reply, {error, relay_mismatch}, State#state{stats = Stats}}
    end.

%% @private
%% @doc Apply the upgrade and update state based on result.
-spec apply_upgrade(binary(), #relay_info{}, term(), map(), #state{}) ->
    {reply, term(), #state{}}.
apply_upgrade(PeerId, RelayInfo, DirectConn, Stats, State) ->
    case do_upgrade(PeerId, RelayInfo, DirectConn) of
        ok ->
            ?LOG_INFO("Connection upgraded to direct for peer ~p", [PeerId]),
            FinalStats = increment_stat(upgrades_succeeded, Stats),
            NewRelays = maps:remove(PeerId, State#state.relays),
            {reply, ok, State#state{relays = NewRelays, stats = FinalStats}};
        {error, Reason} ->
            ?LOG_WARNING("Upgrade failed for peer ~p: ~p", [PeerId, Reason]),
            FinalStats = increment_stat(upgrades_failed, Stats),
            {reply, {error, upgrade_failed}, State#state{stats = FinalStats}}
    end.

%% @private
%% @doc Verify the relay connection is the one we registered.
verify_relay(ProvidedConn, RegisteredConn) ->
    ProvidedConn =:= RegisteredConn.

%% @private
%% @doc Perform the actual connection upgrade.
%% Steps:
%% 1. Pause new messages on relay (mark as upgrading)
%% 2. Drain in-flight messages (short wait)
%% 3. Update routing table to use direct connection
%% 4. Close relay connection gracefully
-spec do_upgrade(binary(), #relay_info{}, term()) -> ok | {error, term()}.
do_upgrade(PeerId, RelayInfo, DirectConn) ->
    handle_upgrade_result(catch execute_upgrade_steps(PeerId, RelayInfo, DirectConn)).

%% @private Execute upgrade steps in sequence
execute_upgrade_steps(PeerId, RelayInfo, DirectConn) ->
    %% Step 1: Mark relay as upgrading (stop accepting new messages)
    mark_relay_upgrading(RelayInfo),
    %% Step 2: Short drain period for in-flight messages
    timer:sleep(?DRAIN_TIMEOUT_MS),
    %% Step 3: Update routing to use direct connection
    update_routing(PeerId, DirectConn),
    %% Step 4: Close relay connection gracefully
    close_relay_gracefully(RelayInfo),
    ok.

%% @private Handle upgrade result
handle_upgrade_result({'EXIT', {Reason, Stack}}) ->
    ?LOG_ERROR("Upgrade failed: ~p~n~p", [Reason, Stack]),
    {error, Reason};
handle_upgrade_result({'EXIT', Reason}) ->
    ?LOG_ERROR("Upgrade failed: ~p", [Reason]),
    {error, Reason};
handle_upgrade_result(ok) ->
    ok;
handle_upgrade_result({error, _Reason} = Error) ->
    Error.

%% @private
%% @doc Mark relay connection as upgrading (no new messages).
mark_relay_upgrading(#relay_info{relay_conn = Conn}) ->
    %% In a full implementation, this would signal the connection
    %% to stop accepting new outbound messages
    ?LOG_DEBUG("Marking relay ~p as upgrading", [Conn]),
    ok.

%% @private
%% @doc Update routing table to use direct connection.
update_routing(PeerId, DirectConn) ->
    %% Notify the connection pool about the new direct connection
    case whereis(macula_peer_connection_pool) of
        undefined ->
            ?LOG_DEBUG("Connection pool not running, routing update skipped");
        _Pid ->
            %% Store the direct connection in the pool
            macula_peer_connection_pool:put(PeerId, DirectConn)
    end,

    %% Also update the peer system if running
    case whereis(macula_peer_system) of
        undefined ->
            ok;
        _PeerPid ->
            %% The peer system may need to know about direct connections
            ?LOG_DEBUG("Notifying peer system of direct connection for ~p", [PeerId])
    end,
    ok.

%% @private
%% @doc Close relay connection gracefully.
close_relay_gracefully(#relay_info{relay_conn = Conn}) ->
    %% Attempt graceful close, ignore errors
    ?LOG_DEBUG("Closing relay connection ~p", [Conn]),
    case catch close_connection(Conn) of
        ok -> ok;
        {'EXIT', _} -> ok;
        {error, _} -> ok
    end.

%% @private
%% @doc Close a connection (abstraction for different connection types).
close_connection(Conn) when is_pid(Conn) ->
    gen_server:stop(Conn, normal, 1000);
close_connection(_Conn) ->
    %% QUIC connection handle - would use quicer:close/1
    %% For now, just return ok
    ok.

%% @private
%% @doc Increment a statistics counter.
increment_stat(Key, Stats) ->
    maps:update_with(Key, fun(V) -> V + 1 end, 1, Stats).