ex_stdlib

%%% @doc
%%% A local, decentralized and scalable key-value process storage.
%%%
%%% It allows developers to lookup one or more processes with a given key.
%%% If the registry has unique keys, a key points to 0 or 1 process.
%%% If the registry allows duplicate keys, a single key may point to any
%%% number of processes. In both cases, different keys could identify the
%%% same process.
%%%
%%% Each entry in the registry is associated to the process that has
%%% registered the key. If the process crashes, the keys associated to that
%%% process are automatically removed. All key comparisons in the registry
%%% are done using the match operation (=:=).
%%%
%%% The registry can be used for different purposes, such as name lookups (using
%%% the via option), storing properties, custom dispatching rules, or a pubsub
%%% implementation.
%%%
%%% == Using in via ==
%%%
%%% Once the registry is started with a given name using
%%% start_link/1, it can be used to register and access named
%%% processes using the {via, registry, {registry_name, key}} tuple:
%%%
%%% ```
%%% {ok, _} = registry:start_link([{keys, unique}, {name, my_registry}]),
%%% Name = {via, registry, {my_registry, "agent"}},
%%% {ok, _} = agent:start_link(fun() -> 0 end, [{name, Name}]),
%%% 0 = agent:get(Name, fun(State) -> State end).
%%% '''
%%%
%%% == Using as a dispatcher ==
%%%
%%% Registry has a dispatch mechanism that allows developers to implement custom
%%% dispatch logic triggered from the caller. For example:
%%%
%%% ```
%%% {ok, _} = registry:start_link([{keys, duplicate}, {name, dispatcher_test}]),
%%% {ok, _} = registry:register(dispatcher_test, "hello", {io, format}),
%%% registry:dispatch(dispatcher_test, "hello", fun(Entries) ->
%%%     [apply(M, F, [Pid, "Hello ~p~n"]) || {Pid, {M, F}} <- Entries]
%%% end).
%%% '''
%%%
%%% @end
-module(registry).

%% API
-export([start_link/1]).
-export([register/3, unregister/2]).
-export([lookup/2, keys/2, values/3]).
-export([dispatch/3, dispatch/4]).
-export([update_value/3]).
-export([via_register/3, via_unregister/3, via_whereis/2, via_send/4]).

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

%% Types
-type keys() :: unique | duplicate.
-type key() :: term().
-type value() :: term().
-type registry() :: atom().
-type entry() :: {pid(), value()}.
-type dispatcher() :: fun(([entry()]) -> term()).

-export_type([keys/0, key/0, value/0, registry/0, entry/0]).

-record(state, {
    keys :: keys(),
    key_table :: ets:tid(),
    pid_table :: ets:tid()
}).

%%% @doc
%%% Starts a registry with the given options.
%%%
%%% The supported options are:
%%%
%%% - {keys, Keys} - either unique or duplicate. Defaults to unique.
%%% - {name, Name} - the name of the registry. Required.
%%%
%%% If keys is unique, a key can only be registered by one process.
%%% If keys is duplicate, multiple processes can register under the same key.
%%%
%%% Returns {ok, Pid} on success.
%%% @end
-spec start_link([{keys, keys()} | {name, atom()}]) -> {ok, pid()} | {error, term()}.
start_link(Options) ->
    Name = proplists:get_value(name, Options),
    case Name of
        undefined ->
            {error, {badarg, "name option is required"}};
        _ ->
            gen_server:start_link({local, Name}, ?MODULE, Options, [])
    end.

%%% @doc
%%% Registers the current process under the given key with an associated value.
%%%
%%% Returns {ok, Owner} on success where Owner is the registry process PID.
%%% Returns {error, {already_registered, Pid}} if the key is already taken
%%% in a unique registry.
%%% @end
-spec register(registry(), key(), value()) -> {ok, pid()} | {error, term()}.
register(Registry, Key, Value) ->
    gen_server:call(Registry, {register, self(), Key, Value}).

%%% @doc
%%% Unregisters the current process from the given key.
%%%
%%% Always returns ok.
%%% @end
-spec unregister(registry(), key()) -> ok.
unregister(Registry, Key) ->
    gen_server:call(Registry, {unregister, self(), Key}).

%%% @doc
%%% Looks up processes registered under the given key.
%%%
%%% Returns a list of {Pid, Value} tuples. For unique registries,
%%% the list will have at most one element.
%%% @end
-spec lookup(registry(), key()) -> [entry()].
lookup(Registry, Key) ->
    try
        gen_server:call(Registry, {lookup, Key})
    catch
        exit:{noproc, _} -> []
    end.

%%% @doc
%%% Returns all keys registered by the given process.
%%% @end
-spec keys(registry(), pid()) -> [key()].
keys(Registry, Pid) ->
    try
        gen_server:call(Registry, {keys, Pid})
    catch
        exit:{noproc, _} -> []
    end.

%%% @doc
%%% Returns the values registered by the given process under the given key.
%%%
%%% For unique registries, returns a list with at most one element.
%%% For duplicate registries, returns a list with zero or more elements.
%%% @end
-spec values(registry(), key(), pid()) -> [value()].
values(Registry, Key, Pid) ->
    try
        gen_server:call(Registry, {values, Key, Pid})
    catch
        exit:{noproc, _} -> []
    end.

%%% @doc
%%% Invokes the callback with all entries under the given key.
%%%
%%% The callback receives a list of {Pid, Value} tuples. If there are
%%% no entries for the given key, the callback is not invoked.
%%% @end
-spec dispatch(registry(), key(), dispatcher()) -> ok.
dispatch(Registry, Key, Callback) ->
    dispatch(Registry, Key, Callback, []).

-spec dispatch(registry(), key(), dispatcher(), list()) -> ok.
dispatch(Registry, Key, Callback, _Options) ->
    case lookup(Registry, Key) of
        [] ->
            ok;
        Entries ->
            Callback(Entries),
            ok
    end.

%%% @doc
%%% Updates the value for the given key for the current process.
%%%
%%% Returns {NewValue, OldValue} on success or error if there is no
%%% such key registered by the current process.
%%%
%%% Only works with unique registries.
%%% @end
-spec update_value(registry(), key(), fun((value()) -> value())) ->
    {value(), value()} | error.
update_value(Registry, Key, Callback) ->
    gen_server:call(Registry, {update_value, self(), Key, Callback}).

%%% @doc
%%% Via callback for process registration.
%%% Used internally by OTP when using {via, registry, {RegistryName, Key}} tuples.
%%% @end
-spec via_register(atom(), term(), pid()) -> yes | no.
via_register(RegistryName, {Key, Value}, _Pid) ->
    case register(RegistryName, Key, Value) of
        {ok, _} -> yes;
        {error, _} -> no
    end;
via_register(RegistryName, Key, _Pid) ->
    via_register(RegistryName, {Key, nil}, _Pid).

%%% @doc
%%% Via callback for process unregistration.
%%% @end
-spec via_unregister(atom(), term(), pid()) -> ok.
via_unregister(RegistryName, {Key, _Value}, _Pid) ->
    unregister(RegistryName, Key);
via_unregister(RegistryName, Key, _Pid) ->
    unregister(RegistryName, Key).

%%% @doc
%%% Via callback for process lookup.
%%% @end
-spec via_whereis(atom(), term()) -> pid() | undefined.
via_whereis(RegistryName, {Key, _Value}) ->
    via_whereis(RegistryName, Key);
via_whereis(RegistryName, Key) ->
    case lookup(RegistryName, Key) of
        [{Pid, _}] -> Pid;
        _ -> undefined
    end.

%%% @doc
%%% Via callback for sending messages.
%%% @end
-spec via_send(atom(), term(), term(), term()) -> pid().
via_send(RegistryName, Name, Msg, _SendOpts) ->
    case via_whereis(RegistryName, Name) of
        undefined ->
            exit({badarg, {RegistryName, Name}});
        Pid ->
            Pid ! Msg,
            Pid
    end.

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

%%% @private
init(Options) ->
    Keys = proplists:get_value(keys, Options, unique),

    case Keys of
        unique -> ok;
        duplicate -> ok;
        _ -> error({badarg, "keys must be unique or duplicate"})
    end,

    % Create ETS tables for key->pid and pid->keys mappings
    KeyTable = ets:new(registry_keys, [
        case Keys of unique -> set; duplicate -> bag end,
        protected,
        {read_concurrency, true}
    ]),

    PidTable = ets:new(registry_pids, [
        bag,
        protected,
        {read_concurrency, true}
    ]),

    {ok, #state{
        keys = Keys,
        key_table = KeyTable,
        pid_table = PidTable
    }}.

%%% @private
handle_call({register, Pid, Key, Value}, _From, State) ->
    #state{keys = Keys, key_table = KeyTable, pid_table = PidTable} = State,

    case Keys of
        unique ->
            case ets:lookup(KeyTable, Key) of
                [] ->
                    % Key not taken, register it
                    monitor(process, Pid),
                    ets:insert(KeyTable, {Key, {Pid, Value}}),
                    ets:insert(PidTable, {Pid, Key}),
                    {reply, {ok, self()}, State};
                [{Key, {ExistingPid, _}}] ->
                    {reply, {error, {already_registered, ExistingPid}}, State}
            end;
        duplicate ->
            % Always allow registration in duplicate mode
            monitor(process, Pid),
            ets:insert(KeyTable, {Key, {Pid, Value}}),
            ets:insert(PidTable, {Pid, Key}),
            {reply, {ok, self()}, State}
    end;

handle_call({unregister, Pid, Key}, _From, State) ->
    #state{key_table = KeyTable, pid_table = PidTable} = State,

    % Remove from both tables
    ets:match_delete(KeyTable, {Key, {Pid, '_'}}),
    ets:delete_object(PidTable, {Pid, Key}),

    % Check if this process has any more keys, if not, demonitor
    case ets:lookup(PidTable, Pid) of
        [] ->
            demonitor_process(Pid);
        _ ->
            ok
    end,

    {reply, ok, State};

handle_call({lookup, Key}, _From, State) ->
    #state{key_table = KeyTable} = State,
    Entries = ets:lookup(KeyTable, Key),
    Result = [{Pid, Value} || {_Key, {Pid, Value}} <- Entries],
    {reply, Result, State};

handle_call({keys, Pid}, _From, State) ->
    #state{pid_table = PidTable} = State,
    Keys = [Key || {_Pid, Key} <- ets:lookup(PidTable, Pid)],
    {reply, Keys, State};

handle_call({values, Key, Pid}, _From, State) ->
    #state{key_table = KeyTable} = State,
    Entries = ets:lookup(KeyTable, Key),
    Values = [Value || {_Key, {P, Value}} <- Entries, P =:= Pid],
    {reply, Values, State};

handle_call({update_value, Pid, Key, Callback}, _From, State) ->
    #state{keys = Keys, key_table = KeyTable} = State,

    case Keys of
        unique ->
            case ets:lookup(KeyTable, Key) of
                [{Key, {Pid, OldValue}}] ->
                    NewValue = Callback(OldValue),
                    ets:insert(KeyTable, {Key, {Pid, NewValue}}),
                    {reply, {NewValue, OldValue}, State};
                _ ->
                    {reply, error, State}
            end;
        duplicate ->
            {reply, {error, "update_value not supported for duplicate registries"}, State}
    end.

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

%%% @private
handle_info({'DOWN', _MonitorRef, process, Pid, _Reason}, State) ->
    #state{key_table = KeyTable, pid_table = PidTable} = State,

    % Get all keys for this process
    Keys = [Key || {_Pid, Key} <- ets:lookup(PidTable, Pid)],

    % Remove all entries for this process
    [ets:match_delete(KeyTable, {Key, {Pid, '_'}}) || Key <- Keys],
    ets:delete(PidTable, Pid),

    {noreply, State};

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

%%% @private
terminate(_Reason, #state{key_table = KeyTable, pid_table = PidTable}) ->
    ets:delete(KeyTable),
    ets:delete(PidTable),
    ok.

%%% @private
code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

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

%% @private
demonitor_process(_Pid) ->
    % In this simplified implementation, we don't track individual monitor refs
    % The process monitor will be automatically cleaned up when the process dies
    ok.