regis

-module(regis_pstore_gb_trees).

-export([ new/0
        , register/4
        , unregister/3
        , unregister_all/2
        , whereis_job/2
        , jobs/2
        , num_registered_pids/1
        ]).

-record(pstore, {proc_store, proc_pids}).
-opaque pstore() :: #pstore{}.
-export_type([pstore/0]).

new() ->
    JobsId = gb_trees:empty(),
    PidsId = gb_trees:empty(),
    {ok, _State=#pstore{proc_store=JobsId, proc_pids=PidsId}}. 

-spec register(pstore(), term(), term(), pid()) -> {ok, pstore()} | {error, exists_proc_store, none_proc_pids} | {error, none_proc_store, exists_proc_pids}.
register(State=#pstore{proc_store=JobsId, proc_pids=PidsId}, JobMeta={_App, _CallBack, _ProcId}, VNodeId, Pid) when is_pid(Pid) -> 
    register_( State
             , JobMeta
             , _VnodeItem={VNodeId, Pid}
             , { gb_trees:lookup(JobMeta, JobsId), gb_trees:lookup(Pid, PidsId) }
             ).

register_(_State=#pstore{proc_store=JobsId, proc_pids=PidsId}, JobMeta, VnodeItem={_VNodeId, Pid}, {none, none}) ->
    {ok, _State1=#pstore{ proc_store=gb_trees:insert(JobMeta, [VnodeItem], JobsId)
                       , proc_pids=gb_trees:insert(Pid, [JobMeta], PidsId)
                       }
    };
register_(_State=#pstore{proc_store=JobsId, proc_pids=PidsId}, JobMeta, VnodeItem={_VNodeId, Pid}, {{value, VnodeItems}, {value, JobMetas}}) ->
    VnodeItem1 = [VnodeItem] ++ [ VI || VI <- VnodeItems, VI/=VnodeItem],
    JobMeta1   = [JobMeta]   ++ [ PI || PI <- JobMetas,   PI/=JobMeta],
    {ok, _State1=#pstore{ proc_store=gb_trees:update(JobMeta, VnodeItem1, JobsId)
                        , proc_pids=gb_trees:update(Pid, JobMeta1, PidsId)
                        }
    };
register_(_State, _JobMeta, _VnodeItem, {{value, _}, none}) ->
    {error, exists_proc_store, none_proc_pids};
register_(_State, _JobMeta, _VnodeItem, {none, {value, _}}) ->
    {error, none_proc_store, exists_proc_pids}.

-spec unregister(pstore(), term(), pid()) -> {ok, pstore()} | {error, exists_proc_store, none_proc_pids} | {error, none_proc_store, exists_proc_pids}.
unregister(State=#pstore{proc_store=JobsId, proc_pids=PidsId}, JobMeta={_App, _CallBack, _ProcId}, Pid) when is_pid(Pid) ->
    unregister_( State
               , JobMeta
               , Pid
               , { gb_trees:lookup(JobMeta, JobsId), gb_trees:lookup(Pid, PidsId) }
    ).

unregister_(State, _JobMeta, _Pid, {none, none}) ->
    {ok, State};
unregister_(_State=#pstore{proc_store=JobsId, proc_pids=PidsId}, JobMeta={_App, _CallBack, _ProcId}, Pid, {{value, VnodeItems}, {value, JobMetas}}) ->
    VnodeItemsFiltered = [{VNodeId, P} || {VNodeId, P} <- VnodeItems, P/=Pid],
    JobMetaFiltered = JobMetas -- [ JobMeta ],
    case {VnodeItemsFiltered, JobMetaFiltered} of
        {[], []} ->
            % unlink(Pid),
            {ok, _State1=#pstore{ proc_store=gb_trees:delete(JobMeta, JobsId)
                               , proc_pids=gb_trees:delete(Pid, PidsId)
                               }
            };
        {[], JobItems1} ->
            {ok, _State1=#pstore{ proc_store=gb_trees:delete(JobMeta, JobsId)
                               , proc_pids=gb_trees:update(Pid, JobItems1, PidsId)
                               }
            };
        {VnodeItems1, []} ->
            {ok, _State1=#pstore{ proc_store=gb_trees:update(JobMeta, VnodeItems1, JobsId)
                               , proc_pids=gb_trees:delete(Pid, PidsId)
                               }
            };
        {VnodeItems1, JobItems1} ->
            {ok, _State1=#pstore{ proc_store=gb_trees:update(JobMeta, VnodeItems1, JobsId)
                               , proc_pids=gb_trees:update(Pid, JobItems1, PidsId)
                               }
            }
    end;
unregister_(_State, _JobMeta, _Pid, {{value, _}, none}) ->
    {error, exists_proc_store, none_proc_pids};
unregister_(_State, _JobMeta, _Pid, {none, {value, _}}) ->
    {error, none_proc_store, exists_proc_pids}.


-spec unregister_all(pstore(), pid()) -> [{ok, pstore()} | {error, exists_proc_store, none_proc_pids} | {error, none_proc_store, exists_proc_pids}].
unregister_all(State=#pstore{proc_store=_JobsId, proc_pids=PidsId}, Pid) ->
    case gb_trees:lookup(Pid, PidsId) of
        none -> 
            [{ok, State}];
        {value, JobItems} ->
            [ unregister(State, JobMeta, Pid) || JobMeta <- JobItems ]
    end.

-spec num_registered_pids(pstore()) -> integer().
num_registered_pids(_State=#pstore{proc_store=_JobsId, proc_pids=PidsId}) ->
    gb_trees:size(PidsId).

-spec whereis_job(pstore(), term()) -> pid() | undefined.
whereis_job(_State=#pstore{proc_store=JobsId, proc_pids=_PidsId}, JobMeta) ->
    case jobs(JobsId, JobMeta) of
        [{_, Pid}] -> Pid;
        _ -> undefined
    end.

-spec jobs(pstore() | atom(), term()) -> [{term(), pid()}].
jobs(_State=#pstore{proc_store=JobsId, proc_pids=_PidsId}, JobMeta) ->
    jobs(JobsId, JobMeta);
jobs(JobsId, JobMeta) ->
    case gb_trees:lookup(JobsId, JobMeta) of
        none -> [];
        {value, VnodeItems} ->
            [{VNodeId, Pid} || {VNodeId, Pid} <- VnodeItems]
    end.