macula

%% @doc Local client for in-VM workloads to connect to macula_gateway
%%
%% This module provides process-to-process communication between workloads
%% running in the same BEAM VM as the Macula platform and the local gateway.
%% Unlike macula_peer which creates QUIC connections, this connects directly
%% to the local macula_gateway process.
%%
%% Architecture:
%%   Phoenix/Elixir App → macula_local_client → macula_gateway
%%                                               ↓ (QUIC)
%%                                          Other Peers
%%
%% @end
-module(macula_local_client).
-behaviour(gen_server).
-behaviour(macula_client_behaviour).

%% API - Connection management
-export([connect/2, connect_local/1, disconnect/1]).
%% API - Pub/Sub
-export([publish/3, publish/4, subscribe/3, unsubscribe/2, discover_subscribers/2]).
%% API - RPC
-export([call/3, call/4, advertise/3, advertise/4, unadvertise/2]).
%% API - Utility
-export([get_node_id/1]).
%% API - Platform Layer (v0.10.0+)
-export([
    register_workload/2,
    get_leader/1,
    subscribe_leader_changes/2,
    propose_crdt_update/3,
    propose_crdt_update/4,
    read_crdt/2
]).
%% Legacy API (kept for backward compatibility)
-export([start_link/1, stop/1, register_procedure/3, unregister_procedure/2]).

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

-record(state, {
    realm :: binary(),
    gateway_pid :: pid() | undefined,
    subscriptions = #{} :: #{reference() => {binary(), fun()}},  % SubRef -> {Topic, Callback}
    registrations = #{} :: #{binary() => fun()},
    event_handler :: pid() | undefined
}).

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

%%------------------------------------------------------------------------------
%% Connection Management
%%------------------------------------------------------------------------------

%% @doc Connect to remote gateway (not supported for local client)
%% For compatibility with macula_client_behaviour
-spec connect(map(), pid()) -> {error, not_supported}.
connect(_Opts, _EventHandler) ->
    {error, not_supported}.

%% @doc Create a local client connection to the gateway
-spec connect_local(map()) -> {ok, pid()} | {error, term()}.
connect_local(Opts) ->
    start_link(Opts).

%% @doc Disconnect the client
-spec disconnect(pid()) -> ok.
disconnect(Pid) ->
    stop(Pid).

%% @doc Start a local client connection to the gateway (legacy API)
-spec start_link(map()) -> {ok, pid()} | {error, term()}.
start_link(Opts) ->
    gen_server:start_link(?MODULE, Opts, []).

%% @doc Stop the local client (legacy API)
-spec stop(pid()) -> ok.
stop(Pid) ->
    gen_server:stop(Pid).

%%------------------------------------------------------------------------------
%% Pub/Sub Operations
%%------------------------------------------------------------------------------

%% @doc Publish an event to a topic
-spec publish(pid(), binary(), map()) -> ok | {error, term()}.
publish(Pid, Topic, Payload) ->
    publish(Pid, Topic, Payload, #{}).

%% @doc Publish an event to a topic with options
-spec publish(pid(), binary(), map(), map()) -> ok | {error, term()}.
publish(Pid, Topic, Payload, Opts) ->
    gen_server:call(Pid, {publish, Topic, Payload, Opts}).

%% @doc Subscribe to a topic
-spec subscribe(pid(), binary(), pid()) -> {ok, reference()} | {error, term()}.
subscribe(Pid, Topic, HandlerPid) ->
    gen_server:call(Pid, {subscribe, Topic, HandlerPid}).

%% @doc Unsubscribe from a topic
-spec unsubscribe(pid(), reference()) -> ok | {error, term()}.
unsubscribe(Pid, SubRef) ->
    gen_server:call(Pid, {unsubscribe, SubRef}).

%% @doc Discover subscribers of a topic via DHT query
-spec discover_subscribers(pid(), binary()) -> {ok, [binary()]} | {error, term()}.
discover_subscribers(Pid, Topic) ->
    gen_server:call(Pid, {discover_subscribers, Topic}).

%%------------------------------------------------------------------------------
%% RPC Operations
%%------------------------------------------------------------------------------

%% @doc Call an RPC procedure with default options
-spec call(pid(), binary(), list()) -> {ok, term()} | {error, term()}.
call(Pid, Procedure, Args) ->
    call(Pid, Procedure, Args, #{}).

%% @doc Call an RPC procedure
-spec call(pid(), binary(), list(), map()) -> {ok, term()} | {error, term()}.
call(Pid, Procedure, Args, Opts) ->
    gen_server:call(Pid, {call, Procedure, Args, Opts}, 30000).

%% @doc Advertise an RPC service with default options
-spec advertise(pid(), binary(), fun()) -> {ok, reference()} | {error, term()}.
advertise(Pid, Procedure, Handler) ->
    advertise(Pid, Procedure, Handler, #{}).

%% @doc Advertise an RPC service with options
-spec advertise(pid(), binary(), fun(), map()) -> {ok, reference()} | {error, term()}.
advertise(Pid, Procedure, Handler, Opts) ->
    gen_server:call(Pid, {advertise, Procedure, Handler, Opts}).

%% @doc Unadvertise an RPC service
-spec unadvertise(pid(), binary()) -> ok | {error, term()}.
unadvertise(Pid, Procedure) ->
    gen_server:call(Pid, {unadvertise, Procedure}).

%% @doc Register an RPC procedure (legacy API, use advertise/3 instead)
-spec register_procedure(pid(), binary(), fun()) -> ok | {error, term()}.
register_procedure(Pid, Procedure, Handler) ->
    advertise(Pid, Procedure, Handler, #{}).

%% @doc Unregister an RPC procedure (legacy API, use unadvertise/2 instead)
-spec unregister_procedure(pid(), binary()) -> ok | {error, term()}.
unregister_procedure(Pid, Procedure) ->
    unadvertise(Pid, Procedure).

%%------------------------------------------------------------------------------
%% Utility Operations
%%------------------------------------------------------------------------------

%% @doc Get the node ID of the local gateway
-spec get_node_id(pid()) -> {ok, binary()} | {error, term()}.
get_node_id(Pid) ->
    gen_server:call(Pid, get_node_id).

%%------------------------------------------------------------------------------
%% Platform Layer Operations (v0.10.0+)
%%------------------------------------------------------------------------------

%% @doc Register this workload with the Platform Layer
-spec register_workload(pid(), map()) -> {ok, map()} | {error, term()}.
register_workload(Pid, Opts) ->
    gen_server:call(Pid, {register_workload, Opts}).

%% @doc Get the current Platform Layer leader node ID
-spec get_leader(pid()) -> {ok, binary()} | {error, no_leader | term()}.
get_leader(Pid) ->
    gen_server:call(Pid, get_leader).

%% @doc Subscribe to Platform Layer leader change notifications
-spec subscribe_leader_changes(pid(), fun((map()) -> ok)) ->
    {ok, reference()} | {error, term()}.
subscribe_leader_changes(Pid, Callback) ->
    gen_server:call(Pid, {subscribe_leader_changes, Callback}).

%% @doc Propose a CRDT update with default options (LWW-Register)
-spec propose_crdt_update(pid(), binary(), term()) -> ok | {error, term()}.
propose_crdt_update(Pid, Key, Value) ->
    propose_crdt_update(Pid, Key, Value, #{crdt_type => lww_register}).

%% @doc Propose a CRDT update with specific type
-spec propose_crdt_update(pid(), binary(), term(), map()) -> ok | {error, term()}.
propose_crdt_update(Pid, Key, Value, Opts) ->
    gen_server:call(Pid, {propose_crdt_update, Key, Value, Opts}).

%% @doc Read the current value of a CRDT-managed state entry
-spec read_crdt(pid(), binary()) -> {ok, term()} | {error, not_found | term()}.
read_crdt(Pid, Key) ->
    gen_server:call(Pid, {read_crdt, Key}).

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

init(Opts) ->
    Realm = maps:get(realm, Opts, <<"default">>),
    EventHandler = maps:get(event_handler, Opts, self()),

    io:format("[LocalClient] Initializing local client for realm ~s~n", [Realm]),

    %% Find the local gateway process
    case find_gateway() of
        {ok, GatewayPid} ->
            io:format("[LocalClient] Connected to local gateway: ~p~n", [GatewayPid]),
            monitor(process, GatewayPid),
            State = #state{
                realm = Realm,
                gateway_pid = GatewayPid,
                event_handler = EventHandler
            },
            {ok, State};
        {error, Reason} ->
            io:format("[LocalClient] Failed to find gateway: ~p~n", [Reason]),
            {stop, {gateway_not_found, Reason}}
    end.

handle_call({publish, Topic, Payload, Opts}, _From, State) ->
    #state{gateway_pid = Gateway, realm = Realm} = State,

    %% Send publish request to gateway via gen_server:call
    %% Note: Opts are currently ignored for local clients
    _ = Opts,
    Result = case gen_server:call(Gateway, {local_publish, Realm, Topic, Payload}) of
        ok -> ok;
        {error, _} = Error -> Error
    end,
    {reply, Result, State};

handle_call({subscribe, Topic, Callback}, _From, State) ->
    #state{gateway_pid = Gateway, realm = Realm, subscriptions = Subs} = State,

    %% Subscribe via local gateway, passing our PID (not the callback)
    %% Gateway will send pubsub events to us, and we'll invoke the callback
    case gen_server:call(Gateway, {local_subscribe, Realm, Topic, self()}) of
        {ok, SubRef} ->
            %% Store callback for later invocation
            NewSubs = maps:put(SubRef, {Topic, Callback}, Subs),
            {reply, {ok, SubRef}, State#state{subscriptions = NewSubs}};
        {error, _} = Error ->
            {reply, Error, State}
    end;

handle_call({unsubscribe, SubRef}, _From, State) ->
    #state{gateway_pid = Gateway, subscriptions = Subs} = State,

    case gen_server:call(Gateway, {local_unsubscribe, SubRef}) of
        ok ->
            %% Remove from local callback tracking
            NewSubs = maps:remove(SubRef, Subs),
            {reply, ok, State#state{subscriptions = NewSubs}};
        {error, _} = Error ->
            {reply, Error, State}
    end;

handle_call({discover_subscribers, Topic}, _From, State) ->
    #state{gateway_pid = Gateway, realm = Realm} = State,

    %% Forward DHT query to gateway
    Result = gen_server:call(Gateway, {local_discover_subscribers, Realm, Topic}),
    {reply, Result, State};

handle_call({call, Procedure, Args, Opts}, _From, State) ->
    #state{gateway_pid = Gateway, realm = Realm} = State,

    %% Route RPC call through local gateway
    Result = gen_server:call(Gateway, {local_rpc_call, Realm, Procedure, Args, Opts}, 30000),
    {reply, Result, State};

handle_call({advertise, Procedure, Handler, Opts}, _From, State) ->
    #state{gateway_pid = Gateway, realm = Realm, registrations = Regs} = State,

    %% Forward advertise request to gateway
    case gen_server:call(Gateway, {local_advertise, Realm, Procedure, Handler, Opts}) of
        {ok, Ref} ->
            NewRegs = maps:put(Procedure, Handler, Regs),
            {reply, {ok, Ref}, State#state{registrations = NewRegs}};
        {error, _} = Error ->
            {reply, Error, State}
    end;

handle_call({unadvertise, Procedure}, _From, State) ->
    #state{gateway_pid = Gateway, registrations = Regs} = State,

    %% Forward unadvertise request to gateway
    case gen_server:call(Gateway, {local_unadvertise, Procedure}) of
        ok ->
            NewRegs = maps:remove(Procedure, Regs),
            {reply, ok, State#state{registrations = NewRegs}};
        {error, _} = Error ->
            {reply, Error, State}
    end;

handle_call({register, Procedure, Handler}, _From, State) ->
    #state{gateway_pid = Gateway, realm = Realm, registrations = Regs} = State,

    case gen_server:call(Gateway, {local_register_procedure, Realm, Procedure, Handler}) of
        ok ->
            NewRegs = maps:put(Procedure, Handler, Regs),
            {reply, ok, State#state{registrations = NewRegs}};
        {error, _} = Error ->
            {reply, Error, State}
    end;

handle_call({unregister, Procedure}, _From, State) ->
    #state{gateway_pid = Gateway, registrations = Regs} = State,

    case gen_server:call(Gateway, {local_unregister_procedure, Procedure}) of
        ok ->
            NewRegs = maps:remove(Procedure, Regs),
            {reply, ok, State#state{registrations = NewRegs}};
        {error, _} = Error ->
            {reply, Error, State}
    end;

handle_call(get_node_id, _From, State) ->
    #state{gateway_pid = Gateway} = State,
    Result = gen_server:call(Gateway, local_get_node_id),
    {reply, Result, State};

%%------------------------------------------------------------------------------
%% Platform Layer Handlers (v0.10.0+)
%%------------------------------------------------------------------------------

handle_call({register_workload, Opts}, _From, State) ->
    %% Query Platform Layer for current state
    Leader = macula_leader_election:get_leader(),
    {ok, Members} = macula_leader_election:get_members(),

    WorkloadName = maps:get(workload_name, Opts, <<"unknown">>),
    Capabilities = maps:get(capabilities, Opts, []),

    io:format("[LocalClient] Registered workload ~s with capabilities ~p~n",
              [WorkloadName, Capabilities]),

    %% Return platform info
    Info = #{
        leader_node => Leader,
        cluster_size => length(Members),
        platform_version => <<"0.10.0">>
    },
    {reply, {ok, Info}, State};

handle_call(get_leader, _From, State) ->
    %% Query current Raft leader from Platform Layer
    case macula_leader_election:get_leader() of
        undefined ->
            {reply, {error, no_leader}, State};
        Leader ->
            {reply, {ok, Leader}, State}
    end;

handle_call({subscribe_leader_changes, Callback}, _From, State) ->
    %% Subscribe to leader change events from Platform Layer
    %% Generate unique callback ID
    CallbackId = make_ref(),

    %% Create wrapper that transforms boolean to map format
    Wrapper = fun(IsLeader) ->
        Leader = macula_leader_election:get_leader(),
        OldLeader = case IsLeader of
            true -> undefined;
            false -> Leader
        end,
        NewLeader = case IsLeader of
            true -> Leader;
            false -> undefined
        end,
        Change = #{
            old_leader => OldLeader,
            new_leader => NewLeader
        },
        Callback(Change)
    end,

    %% Register callback with leader election
    ok = macula_leader_election:register_callback(CallbackId, Wrapper),
    {reply, {ok, CallbackId}, State};

handle_call({propose_crdt_update, Key, Value, Opts}, _From, State) ->
    CrdtType = maps:get(crdt_type, Opts, lww_register),

    %% Propose update to CRDT system
    %% For v0.10.0, store in ETS table (simple implementation)
    %% TODO: Proper CRDT replication in future version
    case CrdtType of
        lww_register ->
            %% LWW-Register: store value with timestamp
            Timestamp = erlang:system_time(millisecond),
            TableName = macula_crdt_storage,

            %% Ensure table exists
            case ets:whereis(TableName) of
                undefined ->
                    try
                        ets:new(TableName, [named_table, public, set])
                    catch
                        error:badarg -> ok  % Table already exists (race condition)
                    end;
                _ -> ok
            end,

            %% Store or update value
            ets:insert(TableName, {Key, {Value, Timestamp}}),
            {reply, ok, State};
        _ ->
            %% Other CRDT types can be added later
            {reply, {error, {unsupported_crdt_type, CrdtType}}, State}
    end;

handle_call({read_crdt, Key}, _From, State) ->
    %% Read current value from CRDT system
    TableName = macula_crdt_storage,

    case ets:whereis(TableName) of
        undefined ->
            {reply, {error, not_found}, State};
        _ ->
            case ets:lookup(TableName, Key) of
                [{Key, {Value, _Timestamp}}] ->
                    {reply, {ok, Value}, State};
                [] ->
                    {reply, {error, not_found}, State}
            end
    end.

%% Async publish - fire-and-forget from caller's perspective
handle_cast({publish_async, Topic, Payload, Opts}, State) ->
    #state{gateway_pid = Gateway, realm = Realm} = State,

    %% Send publish request to gateway via gen_server:cast (async)
    %% Note: Opts are currently ignored for local clients
    _ = Opts,
    gen_server:cast(Gateway, {local_publish_async, Realm, Topic, Payload}),
    {noreply, State};

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

%% @doc Handle pubsub events from gateway
%% Gateway sends messages in format: {publish, Topic, Payload}
handle_info({publish, Topic, Payload}, State) ->
    #state{subscriptions = Subs} = State,

    io:format("[LocalClient] Received publish message for topic ~s~n", [Topic]),
    io:format("[LocalClient] Current subscriptions: ~p~n", [Subs]),

    %% Find all callbacks for subscriptions matching this topic
    %% (since we store SubRef -> {Topic, Callback})
    MatchCount = maps:fold(fun(SubRef, {SubTopic, Callback}, Acc) ->
        io:format("[LocalClient] Checking SubRef=~p, SubTopic=~p against Topic=~p~n",
                  [SubRef, SubTopic, Topic]),
        case SubTopic of
            Topic ->
                %% Topic matches exactly, invoke callback
                io:format("[LocalClient] MATCH FOUND! Invoking callback for topic ~s~n", [Topic]),
                try
                    Callback(Payload),
                    io:format("[LocalClient] Callback completed successfully for topic ~s~n", [Topic])
                catch
                    Class:Reason:Stacktrace ->
                        io:format("[LocalClient] Callback error for topic ~s: ~p:~p~n~p~n",
                                  [Topic, Class, Reason, Stacktrace])
                end,
                Acc + 1;
            _ ->
                %% Different topic, skip
                io:format("[LocalClient] No match: ~p =/= ~p~n", [SubTopic, Topic]),
                Acc
        end
    end, 0, Subs),

    io:format("[LocalClient] Total callbacks invoked for topic ~s: ~p~n", [Topic, MatchCount]),

    {noreply, State};

handle_info({'DOWN', _Ref, process, GatewayPid, Reason}, #state{gateway_pid = GatewayPid} = State) ->
    io:format("[LocalClient] Gateway down: ~p. Attempting reconnect...~n", [Reason]),
    %% Try to reconnect to gateway
    case find_gateway() of
        {ok, NewGateway} ->
            io:format("[LocalClient] Reconnected to gateway: ~p~n", [NewGateway]),
            monitor(process, NewGateway),
            {noreply, State#state{gateway_pid = NewGateway}};
        {error, _} ->
            io:format("[LocalClient] Gateway not available, stopping~n"),
            {stop, gateway_unavailable, State}
    end;

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

terminate(_Reason, #state{subscriptions = Subs, registrations = Regs, gateway_pid = Gateway}) ->
    %% Clean up subscriptions (keys are SubRefs now)
    maps:foreach(fun(SubRef, _TopicCallback) ->
        gen_server:call(Gateway, {local_unsubscribe, SubRef})
    end, Subs),

    %% Clean up registrations (unadvertise all procedures)
    maps:foreach(fun(Procedure, _) ->
        gen_server:call(Gateway, {local_unadvertise, Procedure})
    end, Regs),

    ok.

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

%% @doc Find the local gateway process via whereis
-spec find_gateway() -> {ok, pid()} | {error, not_found}.
find_gateway() ->
    case whereis(macula_gateway) of
        Pid when is_pid(Pid) ->
            {ok, Pid};
        undefined ->
            {error, not_found}
    end.