%%%-------------------------------------------------------------------
%%% @doc Relay-routed Erlang distribution.
%%%
%%% When MACULA_DIST_MODE=relay, Erlang distribution connections
%%% are tunneled through the Macula relay mesh instead of direct QUIC.
%%% This enables distribution across firewalls and NATs — nodes only
%%% need outbound connectivity to a relay.
%%%
%%% The tunnel works by:
%%% 1. Node A publishes a DIST_CONNECT request to the mesh
%%% 2. The relay forwards it to Node B's handler
%%% 3. Node B opens a distribution stream back through the relay
%%% 4. Raw ETF bytes flow through the relay tunnel
%%%
%%% This is EXPERIMENTAL. Supports Pid ! Msg, gen_server:call,
%%% pg groups, process monitoring. Does NOT guarantee Mnesia or
%%% global module compatibility over WAN latency.
%%% @end
%%%-------------------------------------------------------------------
-module(macula_dist_relay).

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

-export([connect/2, accept_dist/2]).
-export([is_relay_mode/0, get_mesh_client/0]).
-export([register_mesh_client/1]).
-export([advertise_dist_accept/0]).

-define(DIST_TOPIC, <<"_dist.connect">>).
-define(DIST_TIMEOUT, 10000).

%% @doc Register a mesh relay client for distribution tunneling.
%% Called by the application that owns the relay client (e.g., hecate_mesh).
-spec register_mesh_client(pid()) -> ok.
register_mesh_client(Pid) when is_pid(Pid) ->
    persistent_term:put(macula_dist_mesh_client, Pid),
    ?LOG_INFO("[dist_relay] Mesh client registered: ~p", [Pid]),
    ok.

%% @doc Check if relay distribution mode is enabled.
-spec is_relay_mode() -> boolean().
is_relay_mode() ->
    case os:getenv("MACULA_DIST_MODE") of
        "relay" -> true;
        _ -> false
    end.

%% @doc Connect to a remote node via relay mesh.
%% Returns {ok, Conn, Stream} where the stream carries ETF bytes
%% through the relay tunnel.
-spec connect(string(), integer()) -> {ok, reference(), reference()} | {error, term()}.
connect(Host, Port) ->
    NodeName = lists:flatten(io_lib:format("~b@~s", [Port, Host])),
    ?LOG_INFO("[dist_relay] Connecting to ~s via mesh", [NodeName]),

    %% Step 1: Find which relay the target node is connected to
    %% For now, use the same relay we're connected to (they share the mesh)
    case get_mesh_client() of
        undefined ->
            {error, no_mesh_connection};
        MeshClient ->
            %% Step 2: Request a distribution tunnel via mesh RPC
            Procedure = <<"_dist.tunnel.", (list_to_binary(NodeName))/binary>>,
            case macula_relay_client:call(MeshClient, Procedure, #{
                <<"from_node">> => atom_to_binary(node()),
                <<"target_node">> => list_to_binary(NodeName)
            }, ?DIST_TIMEOUT) of
                {ok, #{<<"tunnel_id">> := TunnelId}} ->
                    ?LOG_INFO("[dist_relay] Tunnel established: ~s", [TunnelId]),
                    %% Step 3: The tunnel_id identifies a relay-side bridge.
                    %% We use the mesh QUIC connection's stream for ETF bytes.
                    %% For the experimental version, we reuse the existing
                    %% relay_client connection and multiplex dist traffic
                    %% as pub/sub messages on a private topic.
                    create_dist_socket(MeshClient, TunnelId);
                {error, Reason} ->
                    ?LOG_WARNING("[dist_relay] Tunnel request failed: ~p", [Reason]),
                    {error, {tunnel_failed, Reason}}
            end
    end.

%% @doc Accept an incoming distribution connection via relay mesh.
%% Called when a remote node requests a tunnel to us.
-spec accept_dist(binary(), map()) -> {ok, pid()} | {error, term()}.
accept_dist(TunnelId, Opts) ->
    ?LOG_INFO("[dist_relay] Accepting dist tunnel: ~s", [TunnelId]),
    %% The accepting side creates a process that bridges
    %% mesh pub/sub messages to the dist_util handshake
    FromNode = maps:get(<<"from_node">>, Opts, <<>>),
    spawn_link(fun() -> dist_tunnel_process(TunnelId, FromNode) end),
    {ok, TunnelId}.

%% @doc Advertise this node as accepting distribution connections via relay.
%% Registers a `_dist.tunnel.{nodename}` RPC procedure on the mesh.
-spec advertise_dist_accept() -> ok.
advertise_dist_accept() ->
    case get_mesh_client() of
        undefined ->
            ?LOG_WARNING("[dist_relay] Cannot advertise — no mesh client registered"),
            ok;
        Client ->
            NodeName = atom_to_binary(node()),
            Procedure = <<"_dist.tunnel.", NodeName/binary>>,
            Handler = fun(Args) -> handle_tunnel_request(Args) end,
            macula_relay_client:advertise(Client, Procedure, Handler),
            ?LOG_INFO("[dist_relay] Advertised distribution accept: ~s", [Procedure]),
            ok
    end.

%%====================================================================
%% Internal — tunnel negotiation
%%====================================================================

handle_tunnel_request(Args) ->
    FromNode = maps:get(<<"from_node">>, Args, <<>>),
    TunnelId = base64:encode(crypto:strong_rand_bytes(16)),
    ?LOG_INFO("[dist_relay] Tunnel request from ~s → id: ~s", [FromNode, TunnelId]),

    %% Set up the receiving side of the tunnel:
    %% Subscribe to the remote node's outgoing data topic
    %% Publish to the remote node's incoming data topic
    SendTopic = <<"_dist.data.", TunnelId/binary, ".in">>,
    RecvTopic = <<"_dist.data.", TunnelId/binary, ".out">>,

    case get_mesh_client() of
        undefined ->
            {error, <<"no_mesh_client">>};
        Client ->
            %% Subscribe to data from the connecting node
            Self = self(),
            {ok, _SubRef} = macula_relay_client:subscribe(Client, RecvTopic,
                fun(#{payload := Data}) ->
                    Self ! {dist_data, Data}
                end),

            %% Notify net_kernel about the incoming connection
            %% For now, return the tunnel_id so the connecting side can proceed
            {ok, #{<<"tunnel_id">> => TunnelId,
                   <<"send_topic">> => SendTopic,
                   <<"recv_topic">> => RecvTopic}}
    end.

%%====================================================================
%% Internal — mesh client lookup
%%====================================================================

%% Find the mesh relay client. Uses a registered name convention —
%% any application that provides a relay client registers it as
%% macula_dist_mesh_client. No dependency on hecate or any specific app.
get_mesh_client() ->
    case persistent_term:get(macula_dist_mesh_client, undefined) of
        undefined -> undefined;
        Pid when is_pid(Pid) ->
            case is_process_alive(Pid) of
                true -> Pid;
                false -> undefined
            end
    end.

create_dist_socket(MeshClient, TunnelId) ->
    %% Create a pair of linked processes that act as a socket:
    %% - Send: publish ETF bytes to tunnel topic
    %% - Recv: subscribe to tunnel topic, deliver as {quic, Data, ...}
    SendTopic = <<"_dist.data.", TunnelId/binary, ".out">>,
    RecvTopic = <<"_dist.data.", TunnelId/binary, ".in">>,
    Self = self(),

    %% Subscribe to incoming data
    {ok, _SubRef} = macula_relay_client:subscribe(MeshClient, RecvTopic,
        fun(#{payload := Data}) ->
            Self ! {dist_data, Data}
        end),

    {ok, MeshClient, {tunnel, TunnelId, SendTopic, RecvTopic}}.

dist_tunnel_process(_TunnelId, _FromNode) ->
    %% This process lives for the duration of the distribution connection
    receive
        stop -> ok
    end.