%% Copyright 2026 Benoit Chesneau
%%
%% Licensed under the Apache License, Version 2.0 (the "License");
%% you may not use this file except in compliance with the License.
%% You may obtain a copy of the License at
%%
%%     http://www.apache.org/licenses/LICENSE-2.0
%%
%% Unless required by applicable law or agreed to in writing, software
%% distributed under the License is distributed on an "AS IS" BASIS,
%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
%% See the License for the specific language governing permissions and
%% limitations under the License.

%%% @doc Scheduler-affinity router for Python contexts.
%%%
%%% This module provides automatic routing of Python calls to contexts
%%% based on the calling process's scheduler ID. This ensures that
%%% processes on the same scheduler reuse the same context, providing
%%% good cache locality while still enabling N-way parallelism.
%%%
%%% == Architecture ==
%%%
%%% <pre>
%%% Scheduler 1 ---+
%%%                +---> Context 1 (Subinterp/Worker)
%%% Scheduler 2 ---+
%%%                +---> Context 2 (Subinterp/Worker)
%%% Scheduler 3 ---+
%%%                +---> Context 3 (Subinterp/Worker)
%%% ...            |
%%% Scheduler N ---+---> Context N (Subinterp/Worker)
%%% </pre>
%%%
%%% == Usage ==
%%%
%%% <pre>
%%% %% Start the router with default settings
%%% {ok, Contexts} = py_context_router:start(),
%%%
%%% %% Get context for current scheduler (automatic routing)
%%% Ctx = py_context_router:get_context(),
%%% {ok, Result} = py_context:call(Ctx, math, sqrt, [16], #{}),
%%%
%%% %% Or get a specific context by index
%%% Ctx2 = py_context_router:get_context(2),
%%%
%%% %% Bind a specific context to this process
%%% ok = py_context_router:bind_context(Ctx2),
%%% Ctx2 = py_context_router:get_context(), %% Returns bound context
%%%
%%% %% Unbind to return to scheduler-based routing
%%% ok = py_context_router:unbind_context().
%%% </pre>
%%%
%%% @end
-module(py_context_router).

-export([
    start/0,
    start/1,
    stop/0,
    is_started/0,
    get_context/0,
    get_context/1,
    bind_context/1,
    bind_context/2,
    unbind_context/0,
    unbind_context/1,
    num_contexts/0,
    num_contexts/1,
    contexts/0,
    contexts/1,
    %% Pool management
    start_pool/2,
    start_pool/3,
    stop_pool/1,
    pool_started/1,
    %% Pool registration (route module/func to specific pools)
    register_pool/2,
    register_pool/3,
    unregister_pool/1,
    unregister_pool/2,
    lookup_pool/2,
    list_pool_registrations/0,
    init_pool_registry/0
]).

%% Persistent term keys (pool-aware)
-define(POOL_SIZE_KEY(Pool), {?MODULE, pool_size, Pool}).
-define(POOL_CONTEXT_KEY(Pool, N), {?MODULE, pool_context, Pool, N}).
-define(POOL_CONTEXTS_KEY(Pool), {?MODULE, pool_contexts, Pool}).

%% Legacy keys for backward compatibility (default pool)
-define(NUM_CONTEXTS_KEY, {?MODULE, num_contexts}).
-define(CONTEXT_KEY(N), {?MODULE, context, N}).
-define(CONTEXTS_KEY, {?MODULE, all_contexts}).

%% Process dictionary key for bound context (pool-aware)
-define(BOUND_CONTEXT_KEY(Pool), {py_bound_context, Pool}).
-define(BOUND_CONTEXT_KEY, py_bound_context).  %% Legacy for default pool

%% Default pool name
-define(DEFAULT_POOL, default).

%% Pool registry persistent term key
-define(POOL_REGISTRY_KEY, {?MODULE, pool_registry}).

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

-type pool_name() :: default | io | atom().
-type start_opts() :: #{
    contexts => pos_integer(),
    mode => py_context:context_mode()
}.

-export_type([start_opts/0, pool_name/0]).

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

%% @doc Start the context router with default settings.
%%
%% Creates one context per scheduler using worker mode.
%%
%% @returns {ok, [Context]} | {error, Reason}
-spec start() -> {ok, [pid()]} | {error, term()}.
start() ->
    start(#{}).

%% @doc Start the context router with options.
%%
%% Options:
%% - `contexts' - Number of contexts to create (default: number of schedulers)
%% - `mode' - Context mode: `worker' or `owngil' (default: `worker')
%%
%% @param Opts Start options
%% @returns {ok, [Context]} | {error, Reason}
-spec start(start_opts()) -> {ok, [pid()]} | {error, term()}.
start(Opts) ->
    NumContexts = maps:get(contexts, Opts, erlang:system_info(schedulers)),
    Mode = maps:get(mode, Opts, worker),
    %% Delegate to start_pool for the default pool
    start_pool(?DEFAULT_POOL, NumContexts, Mode).

%% @doc Stop the context router.
%%
%% Stops the default pool and removes persistent_term entries.
-spec stop() -> ok.
stop() ->
    stop_pool(?DEFAULT_POOL).

%% @doc Check if contexts have been started and are still alive.
%%
%% @returns true if the default pool is running, false otherwise
-spec is_started() -> boolean().
is_started() ->
    pool_started(?DEFAULT_POOL).

%% @doc Get the context for the current process from the default pool.
%%
%% If the process has a bound context, returns that context.
%% Otherwise, selects a context based on the current scheduler ID.
%%
%% @returns Context pid
-spec get_context() -> pid().
get_context() ->
    get_context(?DEFAULT_POOL).

%% @doc Get a context by index or from a named pool.
%%
%% When called with an integer, gets the Nth context from the default pool.
%% When called with an atom, gets a scheduler-affinity context from that pool.
%%
%% @param NOrPool Context index (1 to num_contexts) or pool name
%% @returns Context pid
-spec get_context(pos_integer() | pool_name()) -> pid().
get_context(N) when is_integer(N), N > 0 ->
    %% Legacy: get Nth context from default pool
    persistent_term:get(?POOL_CONTEXT_KEY(?DEFAULT_POOL, N));
get_context(Pool) when is_atom(Pool) ->
    %% Get context from named pool
    %% Auto-bind on first access to ensure consistent context for a process
    case get(?BOUND_CONTEXT_KEY(Pool)) of
        undefined ->
            Ctx = select_by_scheduler(Pool),
            put(?BOUND_CONTEXT_KEY(Pool), Ctx),
            Ctx;
        Ctx ->
            %% Verify bound context is still alive
            case is_process_alive(Ctx) of
                true ->
                    Ctx;
                false ->
                    %% Re-bind to a new context
                    NewCtx = select_by_scheduler(Pool),
                    put(?BOUND_CONTEXT_KEY(Pool), NewCtx),
                    NewCtx
            end
    end.

%% @doc Bind a context to the current process for the default pool.
%%
%% After binding, `get_context/0' will always return this context
%% instead of selecting by scheduler.
%%
%% @param Ctx Context pid to bind
%% @returns ok
-spec bind_context(pid()) -> ok.
bind_context(Ctx) when is_pid(Ctx) ->
    bind_context(?DEFAULT_POOL, Ctx).

%% @doc Bind a context to the current process for a specific pool.
%%
%% @param Pool Pool name
%% @param Ctx Context pid to bind
%% @returns ok
-spec bind_context(pool_name(), pid()) -> ok.
bind_context(Pool, Ctx) when is_atom(Pool), is_pid(Ctx) ->
    put(?BOUND_CONTEXT_KEY(Pool), Ctx),
    ok.

%% @doc Unbind the current process's context from the default pool.
%%
%% After unbinding, `get_context/0' will return to scheduler-based
%% selection.
%%
%% @returns ok
-spec unbind_context() -> ok.
unbind_context() ->
    unbind_context(?DEFAULT_POOL).

%% @doc Unbind the current process's context from a specific pool.
%%
%% @param Pool Pool name
%% @returns ok
-spec unbind_context(pool_name()) -> ok.
unbind_context(Pool) when is_atom(Pool) ->
    erase(?BOUND_CONTEXT_KEY(Pool)),
    ok.

%% @doc Get the number of contexts in the default pool.
-spec num_contexts() -> non_neg_integer().
num_contexts() ->
    num_contexts(?DEFAULT_POOL).

%% @doc Get the number of contexts in a pool.
-spec num_contexts(pool_name()) -> non_neg_integer().
num_contexts(Pool) when is_atom(Pool) ->
    persistent_term:get(?POOL_SIZE_KEY(Pool), 0).

%% @doc Get all context pids from the default pool.
-spec contexts() -> [pid()].
contexts() ->
    contexts(?DEFAULT_POOL).

%% @doc Get all context pids from a pool.
-spec contexts(pool_name()) -> [pid()].
contexts(Pool) when is_atom(Pool) ->
    persistent_term:get(?POOL_CONTEXTS_KEY(Pool), []).

%% ============================================================================
%% Pool Management
%% ============================================================================

%% @doc Start a named pool with given size.
%%
%% @param Pool Pool name (default, io, or custom)
%% @param Size Number of contexts in the pool
%% @returns {ok, [Context]} | {error, Reason}
-spec start_pool(pool_name(), pos_integer()) -> {ok, [pid()]} | {error, term()}.
start_pool(Pool, Size) ->
    start_pool(Pool, Size, worker).

%% @doc Start a named pool with given size and mode.
%%
%% @param Pool Pool name (default, io, or custom)
%% @param Size Number of contexts in the pool
%% @param Mode Context mode (worker, owngil)
%% @returns {ok, [Context]} | {error, Reason}
-spec start_pool(pool_name(), pos_integer(), py_context:context_mode()) ->
    {ok, [pid()]} | {error, term()}.
start_pool(Pool, Size, Mode) when is_atom(Pool), is_integer(Size), Size > 0 ->
    %% Check if pool already exists
    case persistent_term:get(?POOL_CONTEXTS_KEY(Pool), undefined) of
        undefined ->
            do_start_pool(Pool, Size, Mode);
        Contexts when is_list(Contexts) ->
            %% Verify at least one context is still alive
            case lists:any(fun is_process_alive/1, Contexts) of
                true ->
                    {ok, Contexts};
                false ->
                    %% All dead - restart
                    stop_pool(Pool),
                    do_start_pool(Pool, Size, Mode)
            end
    end.

%% @private
do_start_pool(Pool, Size, Mode) ->
    %% Ensure supervisor is running
    case whereis(py_context_sup) of
        undefined ->
            case py_context_sup:start_link() of
                {ok, _} -> ok;
                {error, {already_started, _}} -> ok;
                Error -> throw(Error)
            end;
        _ ->
            ok
    end,

    %% Start contexts for this pool
    try
        Contexts = lists:map(
            fun(N) ->
                %% Use pool-unique ID: {Pool, N}
                Id = {Pool, N},
                case py_context_sup:start_context(Id, Mode) of
                    {ok, Pid} ->
                        persistent_term:put(?POOL_CONTEXT_KEY(Pool, N), Pid),
                        Pid;
                    {error, Reason} ->
                        throw({context_start_failed, Pool, N, Reason})
                end
            end,
            lists:seq(1, Size)
        ),

        %% Store pool metadata
        persistent_term:put(?POOL_SIZE_KEY(Pool), Size),
        persistent_term:put(?POOL_CONTEXTS_KEY(Pool), Contexts),

        {ok, Contexts}
    catch
        throw:Err ->
            stop_pool(Pool),
            {error, Err}
    end.

%% @doc Stop a named pool.
%%
%% @param Pool Pool name to stop
%% @returns ok
-spec stop_pool(pool_name()) -> ok.
stop_pool(Pool) when is_atom(Pool) ->
    %% Stop all contexts in the pool
    case persistent_term:get(?POOL_CONTEXTS_KEY(Pool), undefined) of
        undefined ->
            ok;
        Contexts ->
            lists:foreach(
                fun(Ctx) ->
                    try py_context_sup:stop_context(Ctx) catch _:_ -> ok end
                end,
                Contexts
            )
    end,

    %% Clean up persistent terms for this pool
    Size = persistent_term:get(?POOL_SIZE_KEY(Pool), 0),
    lists:foreach(
        fun(N) ->
            try persistent_term:erase(?POOL_CONTEXT_KEY(Pool, N)) catch _:_ -> ok end
        end,
        lists:seq(1, Size)
    ),
    try persistent_term:erase(?POOL_SIZE_KEY(Pool)) catch _:_ -> ok end,
    try persistent_term:erase(?POOL_CONTEXTS_KEY(Pool)) catch _:_ -> ok end,
    ok.

%% @doc Check if a pool has been started and is still alive.
%%
%% @param Pool Pool name to check
%% @returns true if pool is running, false otherwise
-spec pool_started(pool_name()) -> boolean().
pool_started(Pool) when is_atom(Pool) ->
    case persistent_term:get(?POOL_SIZE_KEY(Pool), 0) of
        0 -> false;
        _N ->
            %% Verify at least one context is actually alive
            case persistent_term:get(?POOL_CONTEXT_KEY(Pool, 1), undefined) of
                undefined -> false;
                Pid when is_pid(Pid) -> is_process_alive(Pid);
                _ -> false
            end
    end.

%% ============================================================================
%% Pool Registration (route module/func to specific pools)
%% ============================================================================

%% @doc Initialize the pool registry.
%%
%% This is called automatically on first registration.
%% Safe to call multiple times - does nothing if already initialized.
-spec init_pool_registry() -> ok.
init_pool_registry() ->
    case persistent_term:get(?POOL_REGISTRY_KEY, undefined) of
        undefined ->
            persistent_term:put(?POOL_REGISTRY_KEY, #{}),
            ok;
        _ ->
            ok
    end.

%% @doc Register a module to use a specific pool for all functions.
%%
%% All calls to `Module:*' will be routed to the specified pool.
%%
%% Example:
%% ```
%% %% Route all 'requests' module calls to io pool
%% py_context_router:register_pool(io, requests).
%% '''
%%
%% @param Pool Pool name to route to
%% @param Module Python module name
%% @returns ok
-spec register_pool(pool_name(), atom()) -> ok.
register_pool(Pool, Module) when is_atom(Pool), is_atom(Module) ->
    Registry = get_registry(),
    NewRegistry = maps:put({Module, '_'}, Pool, Registry),
    persistent_term:put(?POOL_REGISTRY_KEY, NewRegistry),
    ok.

%% @doc Register a specific module/function to use a specific pool.
%%
%% Calls to `Module:Func' will be routed to the specified pool.
%% More specific registrations (module+func) take precedence over
%% module-only registrations.
%%
%% Example:
%% ```
%% %% Route requests.get to io pool
%% py_context_router:register_pool(io, requests, get).
%% '''
%%
%% @param Pool Pool name to route to
%% @param Module Python module name
%% @param Func Python function name
%% @returns ok
-spec register_pool(pool_name(), atom(), atom()) -> ok.
register_pool(Pool, Module, Func) when is_atom(Pool), is_atom(Module), is_atom(Func) ->
    Registry = get_registry(),
    NewRegistry = maps:put({Module, Func}, Pool, Registry),
    persistent_term:put(?POOL_REGISTRY_KEY, NewRegistry),
    ok.

%% @doc Unregister a module from pool routing.
%%
%% Calls will return to using the default pool.
%%
%% @param Module Python module name
%% @returns ok
-spec unregister_pool(atom()) -> ok.
unregister_pool(Module) when is_atom(Module) ->
    Registry = get_registry(),
    NewRegistry = maps:remove({Module, '_'}, Registry),
    persistent_term:put(?POOL_REGISTRY_KEY, NewRegistry),
    ok.

%% @doc Unregister a specific module/function from pool routing.
%%
%% @param Module Python module name
%% @param Func Python function name
%% @returns ok
-spec unregister_pool(atom(), atom()) -> ok.
unregister_pool(Module, Func) when is_atom(Module), is_atom(Func) ->
    Registry = get_registry(),
    NewRegistry = maps:remove({Module, Func}, Registry),
    persistent_term:put(?POOL_REGISTRY_KEY, NewRegistry),
    ok.

%% @doc Look up which pool a module/function is registered to.
%%
%% First checks for a specific module+func registration.
%% If not found, checks for a module-only registration.
%% Returns `default' if no registration found.
%%
%% @param Module Python module name
%% @param Func Python function name
%% @returns Pool name
-spec lookup_pool(atom(), atom()) -> pool_name().
lookup_pool(Module, Func) when is_atom(Module), is_atom(Func) ->
    Registry = get_registry(),
    %% First try specific module+func, then module-only, then default
    case maps:get({Module, Func}, Registry, undefined) of
        undefined ->
            case maps:get({Module, '_'}, Registry, undefined) of
                undefined -> ?DEFAULT_POOL;
                Pool -> Pool
            end;
        Pool ->
            Pool
    end.

%% @doc List all pool registrations.
%%
%% Returns a list of `{{Module, Func}, Pool}' tuples.
%% @returns List of registrations
-spec list_pool_registrations() -> [{{atom(), atom() | '_'}, pool_name()}].
list_pool_registrations() ->
    maps:to_list(get_registry()).

%% @private
%% Get the pool registry map from persistent_term
-spec get_registry() -> #{}.
get_registry() ->
    persistent_term:get(?POOL_REGISTRY_KEY, #{}).

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

%% @private
%% Select context based on scheduler ID using modulo
-spec select_by_scheduler(pool_name()) -> pid().
select_by_scheduler(Pool) ->
    SchedId = erlang:system_info(scheduler_id),
    NumCtx = persistent_term:get(?POOL_SIZE_KEY(Pool)),
    Idx = ((SchedId - 1) rem NumCtx) + 1,
    persistent_term:get(?POOL_CONTEXT_KEY(Pool, Idx)).