%%% @doc
%%% PartitionSupervisor implementation inspired by Elixir's PartitionSupervisor.
%%%
%%% A PartitionSupervisor is a supervisor that manages multiple partitions of the same
%%% child specification. It's designed for high-concurrency scenarios where you want
%%% to distribute work across multiple instances of the same process type.
%%%
%%% Each partition runs as a separate DynamicSupervisor, and work is distributed
%%% across partitions using a hash function on a partition key.
%%%
%%% Examples:
%%% ```
%%% % Start a partition supervisor with 4 partitions
%%% {ok, Pid} = partition_supervisor:start_link([
%%%     {partitions, 4},
%%%     {child_spec, #{id => worker, start => {worker, start_link, []}}}
%%% ]),
%%%
%%% % Start a child in a specific partition
%%% {ok, ChildPid} = partition_supervisor:start_child(Pid, user_123, [arg1, arg2]).
%%% '''
%%% @end
-module(partition_supervisor).

-behaviour(supervisor).

%% API
-export([start_link/1, start_link/2]).
-export([start_child/3, start_child/4]).
-export([terminate_child/2, terminate_child/3]).
-export([which_children/1, which_children/2]).
-export([count_children/1, count_children/2]).
-export([partitions/1, partition_for/2]).
-export([whereis_name/2]).

%% Supervisor callbacks
-export([init/1]).

%% Types
-type partition_key() :: term().
-type partition_index() :: non_neg_integer().
-type partition_options() :: #{
    partitions := pos_integer(),
    child_spec := supervisor:child_spec(),
    hash_function => fun((partition_key()) -> non_neg_integer()),
    name => atom()
}.

-record(state, {
    partitions :: pos_integer(),
    child_spec :: supervisor:child_spec(),
    hash_function :: fun((partition_key()) -> non_neg_integer()),
    partition_pids :: #{partition_index() => pid()}
}).

-export_type([partition_key/0, partition_index/0, partition_options/0]).

%%% @doc
%%% Starts a PartitionSupervisor with the given options.
%%%
%%% Options:
%%% - partitions: Number of partitions to create (required)
%%% - child_spec: Child specification for each partition (required)
%%% - hash_function: Custom hash function (optional, defaults to erlang:phash2/2)
%%% @end
-spec start_link(partition_options()) -> {ok, pid()} | {error, term()}.
start_link(Options) ->
    supervisor:start_link(?MODULE, Options).

%%% @doc
%%% Starts a named PartitionSupervisor with the given options.
%%%
%%% The supervisor will be registered under the given name.
%%% @end
-spec start_link(atom(), partition_options()) -> {ok, pid()} | {error, term()}.
start_link(Name, Options) ->
    supervisor:start_link({local, Name}, ?MODULE, Options#{name => Name}).

%%% @doc
%%% Starts a child in the partition determined by the partition key.
%%%
%%% The partition is selected using the hash function on the partition key.
%%% @end
-spec start_child(pid() | atom(), partition_key(), [term()]) ->
    {ok, pid()} | {ok, pid(), term()} | {error, term()}.
start_child(Supervisor, PartitionKey, Args) ->
    start_child(Supervisor, PartitionKey, Args, []).

%%% @doc
%%% Starts a child in the partition determined by the partition key with options.
%%%
%%% Additional options can be passed to modify the child spec.
%%% @end
-spec start_child(pid() | atom(), partition_key(), [term()], [term()]) ->
    {ok, pid()} | {ok, pid(), term()} | {error, term()}.
start_child(Supervisor, PartitionKey, Args, _Options) ->
    PartitionIndex = partition_for(Supervisor, PartitionKey),
    PartitionPid = get_partition_pid(Supervisor, PartitionIndex),

    % Get the base child spec and modify it with the provided args
    ChildSpec = get_child_spec(Supervisor),
    ModifiedSpec = modify_child_spec_args(ChildSpec, Args),

    dynamic_supervisor:start_child(PartitionPid, ModifiedSpec).

%%% @doc
%%% Terminates a child in the partition determined by the partition key.
%%%
%%% The child is identified by its PID.
%%% @end
-spec terminate_child(pid() | atom(), pid()) -> ok | {error, term()}.
terminate_child(Supervisor, ChildPid) ->
    % Find which partition contains this child
    case find_partition_for_child(Supervisor, ChildPid) of
        {ok, PartitionPid} ->
            dynamic_supervisor:terminate_child(PartitionPid, ChildPid);
        {error, not_found} ->
            {error, not_found}
    end.

%%% @doc
%%% Terminates a child in a specific partition.
%%%
%%% The child is identified by its PID and the partition is specified directly.
%%% @end
-spec terminate_child(pid() | atom(), partition_index(), pid()) -> ok | {error, term()}.
terminate_child(Supervisor, PartitionIndex, ChildPid) ->
    PartitionPid = get_partition_pid(Supervisor, PartitionIndex),
    dynamic_supervisor:terminate_child(PartitionPid, ChildPid).

%%% @doc
%%% Returns information about all children in all partitions.
%%%
%%% Returns a list of {Id, Child, Type, Modules} tuples for all children.
%%% @end
-spec which_children(pid() | atom()) -> [{term(), pid() | undefined, worker | supervisor, [atom()] | dynamic}].
which_children(Supervisor) ->
    PartitionCount = partitions(Supervisor),
    lists:flatten([
        which_children(Supervisor, I) || I <- lists:seq(0, PartitionCount - 1)
    ]).

%%% @doc
%%% Returns information about children in a specific partition.
%%%
%%% Returns a list of {Id, Child, Type, Modules} tuples for children in the partition.
%%% @end
-spec which_children(pid() | atom(), partition_index()) -> [{term(), pid() | undefined, worker | supervisor, [atom()] | dynamic}].
which_children(Supervisor, PartitionIndex) ->
    PartitionPid = get_partition_pid(Supervisor, PartitionIndex),
    dynamic_supervisor:which_children(PartitionPid).

%%% @doc
%%% Returns a count of children in all partitions.
%%%
%%% Returns a property list with counts of active, specs, supervisors, and workers.
%%% @end
-spec count_children(pid() | atom()) -> [{specs | active | supervisors | workers, non_neg_integer()}].
count_children(Supervisor) ->
    PartitionCount = partitions(Supervisor),
    Counts = [count_children(Supervisor, I) || I <- lists:seq(0, PartitionCount - 1)],
    sum_counts(Counts).

%%% @doc
%%% Returns a count of children in a specific partition.
%%%
%%% Returns a property list with counts of active, specs, supervisors, and workers.
%%% @end
-spec count_children(pid() | atom(), partition_index()) -> #{specs := non_neg_integer(), active := non_neg_integer(), supervisors := non_neg_integer(), workers := non_neg_integer()}.
count_children(Supervisor, PartitionIndex) ->
    PartitionPid = get_partition_pid(Supervisor, PartitionIndex),
    dynamic_supervisor:count_children(PartitionPid).

%%% @doc
%%% Returns the number of partitions in the supervisor.
%%% @end
-spec partitions(pid() | atom()) -> pos_integer().
partitions(Supervisor) ->
    {ok, Config} = get_supervisor_config(Supervisor),
    Config#state.partitions.

%%% @doc
%%% Returns the partition index for the given partition key.
%%%
%%% Uses the configured hash function to determine the partition.
%%% @end
-spec partition_for(pid() | atom(), partition_key()) -> partition_index().
partition_for(Supervisor, PartitionKey) ->
    {ok, Config} = get_supervisor_config(Supervisor),
    HashFunction = Config#state.hash_function,
    PartitionCount = Config#state.partitions,
    HashFunction(PartitionKey) rem PartitionCount.

%%% @doc
%%% Returns the PID of a child process by partition key.
%%%
%%% This is useful for direct communication with processes in specific partitions.
%%% @end
-spec whereis_name(pid() | atom(), partition_key()) -> pid() | undefined.
whereis_name(Supervisor, PartitionKey) ->
    PartitionIndex = partition_for(Supervisor, PartitionKey),
    get_partition_pid(Supervisor, PartitionIndex).

%%%=============================================================================
%%% Supervisor callbacks
%%%=============================================================================

%% @private
init(Options) ->
    Partitions = maps:get(partitions, Options),
    ChildSpec = maps:get(child_spec, Options),
    HashFunction = maps:get(hash_function, Options, fun erlang:phash2/1),

    % Store configuration in process dictionary for access
    Config = #state{
        partitions = Partitions,
        child_spec = ChildSpec,
        hash_function = HashFunction,
        partition_pids = #{}
    },
    put(partition_supervisor_config, Config),

    % Create child specifications for each partition
    PartitionSpecs = [
        #{
            id => {partition, I},
            start => {dynamic_supervisor, start_link, [[]]},
            restart => permanent,
            shutdown => infinity,
            type => supervisor,
            modules => [dynamic_supervisor]
        }
        || I <- lists:seq(0, Partitions - 1)
    ],

    SupFlags = #{
        strategy => one_for_one,
        intensity => 10,
        period => 5
    },

    {ok, {SupFlags, PartitionSpecs}}.

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

%% @private
get_supervisor_config(Supervisor) when is_pid(Supervisor) ->
    case erlang:process_info(Supervisor, dictionary) of
        {dictionary, Dict} ->
            case lists:keyfind(partition_supervisor_config, 1, Dict) of
                {partition_supervisor_config, Config} -> {ok, Config};
                false -> {error, config_not_found}
            end;
        undefined -> {error, process_not_found}
    end;
get_supervisor_config(Name) when is_atom(Name) ->
    case whereis(Name) of
        undefined -> {error, supervisor_not_found};
        Pid -> get_supervisor_config(Pid)
    end.

%% @private
get_partition_pid(Supervisor, PartitionIndex) ->
    Children = supervisor:which_children(Supervisor),
    case lists:keyfind({partition, PartitionIndex}, 1, Children) of
        {{partition, PartitionIndex}, Pid, supervisor, _} -> Pid;
        false -> error({partition_not_found, PartitionIndex})
    end.

%% @private
get_child_spec(Supervisor) ->
    {ok, Config} = get_supervisor_config(Supervisor),
    Config#state.child_spec.

%% @private
modify_child_spec_args(ChildSpec, Args) ->
    case ChildSpec of
        #{start := {M, F, ExistingArgs}} = Spec ->
            Spec#{start := {M, F, ExistingArgs ++ Args}};
        {Id, {M, F, ExistingArgs}, Restart, Shutdown, Type, Modules} ->
            {Id, {M, F, ExistingArgs ++ Args}, Restart, Shutdown, Type, Modules}
    end.

%% @private
find_partition_for_child(Supervisor, ChildPid) ->
    PartitionCount = partitions(Supervisor),
    find_partition_for_child_loop(Supervisor, ChildPid, 0, PartitionCount).

%% @private
find_partition_for_child_loop(_Supervisor, _ChildPid, PartitionIndex, PartitionCount)
  when PartitionIndex >= PartitionCount ->
    {error, not_found};
find_partition_for_child_loop(Supervisor, ChildPid, PartitionIndex, PartitionCount) ->
    PartitionPid = get_partition_pid(Supervisor, PartitionIndex),
    Children = dynamic_supervisor:which_children(PartitionPid),
    case lists:any(fun({_, Pid, _, _}) -> Pid =:= ChildPid end, Children) of
        true -> {ok, PartitionPid};
        false -> find_partition_for_child_loop(Supervisor, ChildPid, PartitionIndex + 1, PartitionCount)
    end.

%% @private
sum_counts(CountsList) ->
    InitialCounts = [{active, 0}, {specs, 0}, {supervisors, 0}, {workers, 0}],
    lists:foldl(fun(Counts, Acc) ->
        % Convert map to proplist if needed
        CountsProplist = case Counts of
            #{} = Map -> maps:to_list(Map);
            List when is_list(List) -> List
        end,
        lists:foldl(fun({Key, Value}, AccInner) ->
            OldValue = proplists:get_value(Key, AccInner, 0),
            lists:keystore(Key, 1, AccInner, {Key, OldValue + Value})
        end, Acc, CountsProplist)
    end, InitialCounts, CountsList).