-module(ides_march).

-moduledoc "Kill graph analysis and restart logic for ides.".

-export([
    kill_graph/1,
    should_restart/2,
    affected_siblings/1,
    link_info/1,
    monitor_info/1,
    kill_graph_detail/1,
    intensity_info/1,
    init_analysis/1
]).

-type exit_reason() :: normal | abnormal.
-export_type([exit_reason/0]).

-doc """
Return all PIDs that could cause `TargetPid` to be killed.

This is the set of ancestors unioned with siblings that trigger
cascade restarts under the parent supervisor's strategy:
- `one_for_one` / `simple_one_for_one`: only ancestors (no
  sibling killers)
- `one_for_all`: all siblings are killers
- `rest_for_one`: siblings at positions before the target are
  killers
""".
-spec kill_graph(TargetPid :: pid()) -> {ok, [pid()]} | {error, term()}.
kill_graph(TargetPid) ->
    case ides_family:get_ancestors(TargetPid) of
        {ok, Ancestors} when Ancestors =/= [] ->
            AncestorPids = resolve_ancestor_pids(Ancestors),
            case ides_family:parent_info(TargetPid) of
                {ok, Info} ->
                    KillerSiblings = killer_siblings(Info),
                    LinkKillers = link_killers(TargetPid),
                    MonitorKillers = monitor_killers(TargetPid),
                    {ok,
                        lists:usort(
                            AncestorPids ++
                                KillerSiblings ++
                                LinkKillers ++
                                MonitorKillers
                        )};
                {error, Reason} ->
                    {error, Reason}
            end;
        {ok, []} ->
            {error, no_ancestors};
        {error, Reason} ->
            {error, Reason}
    end.

-doc """
Return whether a terminated child `Pid` would be restarted by
its supervisor.

Rules:
- `permanent`: always restarted
- `transient`: restarted only on `abnormal` exit
- `temporary`: never restarted
""".
-spec should_restart(Pid :: pid(), Reason :: exit_reason()) -> boolean().
should_restart(Pid, ExitReason) ->
    case ides_family:get_ancestors(Pid) of
        {ok, Ancestors} when Ancestors =/= [] ->
            [Parent | _] = Ancestors,
            ParentPid = ides_family:resolve_pid(Parent),
            ChildSpecs = supervisor:which_children(ParentPid),
            {Id, _ChildPid, _Type, _Mods} = lists:keyfind(Pid, 2, ChildSpecs),
            RestartType = ides_family:get_restart_type(ParentPid, Id),
            should_restart_rule(RestartType, ExitReason);
        _ ->
            false
    end.

-doc """
Return the PIDs of siblings that would be killed or restarted
if `TargetPid` dies.

Depends on the parent supervisor's strategy:
- `one_for_one`: `TargetPid` only
- `one_for_all`: all siblings
- `rest_for_one`: `TargetPid` and all siblings after it
- `simple_one_for_one`: `TargetPid` only
""".
-spec affected_siblings(TargetPid :: pid()) -> {ok, [pid()]} | {error, term()}.
affected_siblings(TargetPid) ->
    case ides_family:parent_info(TargetPid) of
        {ok, Info} ->
            Affected = affected_siblings_rule(Info),
            {ok, Affected};
        {error, Reason} ->
            {error, Reason}
    end.

-doc """
Return link information for the given process.

Reports which processes are linked to `Pid` and whether
`Pid` traps exits. Linked processes are potential killers
if `traps_exits` is `false`.
""".
-spec link_info(Pid :: pid()) -> {ok, ides_family:link_info()} | {error, term()}.
link_info(Pid) ->
    case erlang:process_info(Pid, [links, trap_exit, status]) of
        [{links, Links}, {trap_exit, Traps}, {status, _}] ->
            {ok, #{
                links => Links -- [Pid],
                traps_exits => Traps
            }};
        undefined ->
            {error, process_not_alive};
        Other ->
            {error, {unexpected_process_info, Other}}
    end.

-doc """
Return monitor information for the given process.

Reports which processes `Pid` is monitoring and which
processes are monitoring `Pid`. Monitored processes are
potential killers if `Pid` doesn't handle DOWN messages.
""".
-spec monitor_info(Pid :: pid()) -> {ok, ides_family:monitor_info()} | {error, term()}.
monitor_info(Pid) ->
    case erlang:process_info(Pid, [monitors, monitored_by, status]) of
        [{monitors, Monitors}, {monitored_by, MonitoredBy}, {status, _}] ->
            MonitorPids = [MPid || {process, MPid} <- Monitors],
            {ok, #{
                monitors => MonitorPids,
                monitored_by => MonitoredBy
            }};
        undefined ->
            {error, process_not_alive};
        Other ->
            {error, {unexpected_process_info, Other}}
    end.

-doc """
Return the kill graph with each entry tagged by its kill mechanism:
- `ancestor` — supervisor ancestor
- `sibling` — sibling via supervisor strategy
- `link` — linked process (relevant if target doesn't trap exits)
- `monitor` — monitored process (relevant if target doesn't handle DOWN)
""".
-spec kill_graph_detail(TargetPid :: pid()) -> {ok, [ides_family:kill_source()]} | {error, term()}.
kill_graph_detail(TargetPid) ->
    case ides_family:get_ancestors(TargetPid) of
        {ok, Ancestors} when Ancestors =/= [] ->
            AncestorPids = resolve_ancestor_pids(Ancestors),
            case ides_family:parent_info(TargetPid) of
                {ok, Info} ->
                    SiblingKillers = killer_siblings(Info),
                    LinkKillers = link_killers(TargetPid),
                    MonitorKillers = monitor_killers(TargetPid),
                    Tagged =
                        [{ancestor, P} || P <- AncestorPids] ++
                            [{sibling, P} || P <- SiblingKillers] ++
                            [{link, P} || P <- LinkKillers] ++
                            [{monitor, P} || P <- MonitorKillers],
                    {ok, Tagged};
                {error, Reason} ->
                    {error, Reason}
            end;
        {ok, []} ->
            {error, no_ancestors};
        {error, Reason} ->
            {error, Reason}
    end.

-doc """
Return restart intensity information for a supervisor.

Reports the configured MaxR and MaxT policy, plus the current
restart count and remaining budget if available from OTP state.
""".
-spec intensity_info(SupPid :: pid()) -> {ok, ides_family:intensity_info()} | {error, term()}.
intensity_info(SupPid) ->
    try sys:get_state(SupPid) of
        State when is_tuple(State), tuple_size(State) >= 7 ->
            Intensity = element(5, State),
            Period = element(6, State),
            Restarts = element(7, State),
            case {is_integer(Intensity), is_integer(Period), is_list(Restarts)} of
                {true, true, true} ->
                    Count = length(Restarts),
                    {ok, #{
                        max_restarts => Intensity,
                        max_period => Period,
                        current_count => Count,
                        remaining => max(0, Intensity - Count)
                    }};
                _ ->
                    {ok, #{max_restarts => 1, max_period => 5}}
            end;
        _ ->
            {ok, #{max_restarts => 1, max_period => 5}}
    catch
        _:_ ->
            {ok, #{max_restarts => 1, max_period => 5}}
    end.

-doc """
Return init analysis for the parent supervisor of `Pid`.

Walks the parent supervisor's children and computes a worst-case
init failure assessment: how many children could fail init before
exhausting the supervisor's restart intensity budget.
""".
-spec init_analysis(Pid :: pid()) ->
    {ok, ides_family:init_analysis_result()} | {error, term()}.
init_analysis(Pid) ->
    case ides_family:parent_info(Pid) of
        {ok, #{sup_pid := SupPid, sup_strategy := Strategy, child_pids := ChildPids}} ->
            Children = build_children(SupPid, ChildPids),
            {ok, Intensity} = intensity_info(SupPid),
            WorstCase = length([1 || #{counts_against_intensity := true} <- Children]),
            Budget =
                maps:get(max_restarts, Intensity) -
                    maps:get(current_count, Intensity, 0),
            {ok, #{
                supervisor => SupPid,
                sup_strategy => Strategy,
                sup_intensity => Intensity,
                target_pid => Pid,
                total_children => length(Children),
                children => Children,
                worst_case_restarts => WorstCase,
                remaining_budget => Budget
            }};
        {error, Reason} ->
            {error, Reason}
    end.

-spec build_children(SupPid :: pid(), ChildPids :: [{term(), pid()}]) ->
    [ides_family:child_init_info()].
build_children(SupPid, ChildPids) ->
    RawChildren = supervisor:which_children(SupPid),
    [build_child_init_info(SupPid, Id, Pid, RawChildren) || {Id, Pid} <- ChildPids].

-spec build_child_init_info(
    SupPid :: pid(), Id :: term(), Pid :: pid(), RawChildren :: [term()]
) -> ides_family:child_init_info().
build_child_init_info(SupPid, Id, Pid, RawChildren) ->
    RestartType = ides_family:get_restart_type(SupPid, Id),
    Shutdown = get_shutdown(SupPid, Id),
    Phase =
        case lists:keyfind(Id, 1, RawChildren) of
            {Id, Pid, _Type, _Mods} when is_pid(Pid) -> running;
            {Id, restarting, _Type, _Mods} -> restarting;
            _ -> undefined
        end,
    Counts = counts_against_intensity(RestartType),
    #{
        id => Id,
        pid => Pid,
        restart_type => RestartType,
        shutdown => Shutdown,
        phase => Phase,
        counts_against_intensity => Counts
    }.

-spec get_shutdown(SupPid :: pid(), Id :: term()) -> timeout().
get_shutdown(SupPid, Id) ->
    case supervisor:get_childspec(SupPid, Id) of
        {ok, #{shutdown := Shutdown}} -> Shutdown;
        {ok, {Id, _StartFunc, _RestartType, Shutdown, _Type, _Modules}} -> Shutdown;
        _ -> 5000
    end.

-spec counts_against_intensity(RestartType :: ides_family:child_restart_type()) -> boolean().
counts_against_intensity(permanent) -> true;
counts_against_intensity(transient) -> true;
counts_against_intensity(temporary) -> false.

%% --- Internal ---

-spec resolve_ancestor_pids(Ancestors :: [term()]) -> [pid()].
resolve_ancestor_pids(Ancestors) ->
    lists:filtermap(
        fun(P) ->
            try ides_family:resolve_pid(P) of
                Pid -> {true, Pid}
            catch
                _:_ -> false
            end
        end,
        Ancestors
    ).

-spec killer_siblings(Info :: ides_family:parent_info()) -> [pid()].
killer_siblings(#{sup_strategy := one_for_all, child_pids := Children, target_position := _Pos}) ->
    pids(Children);
killer_siblings(#{sup_strategy := rest_for_one, child_pids := Children, target_position := Pos}) ->
    pids(child_prefix(Children, Pos - 1));
killer_siblings(#{sup_strategy := _, child_pids := _Children, target_position := _Pos}) ->
    [].

-spec should_restart_rule(
    RestartType :: ides_family:child_restart_type(), ExitReason :: exit_reason()
) -> boolean().
should_restart_rule(permanent, _) -> true;
should_restart_rule(transient, abnormal) -> true;
should_restart_rule(transient, normal) -> false;
should_restart_rule(temporary, _) -> false.

-spec affected_siblings_rule(Info :: ides_family:parent_info()) -> [pid()].
affected_siblings_rule(#{
    sup_strategy := one_for_one, child_pids := Children, target_position := Pos
}) ->
    pids(child_sublist(Children, Pos, 1));
affected_siblings_rule(#{
    sup_strategy := one_for_all, child_pids := Children, target_position := _Pos
}) ->
    pids(Children);
affected_siblings_rule(#{
    sup_strategy := rest_for_one, child_pids := Children, target_position := Pos
}) ->
    pids(child_suffix(Children, Pos - 1));
affected_siblings_rule(#{
    sup_strategy := simple_one_for_one, child_pids := Children, target_position := Pos
}) ->
    pids(child_sublist(Children, Pos, 1)).

-spec link_killers(Pid :: pid()) -> [pid()].
link_killers(Pid) ->
    case link_info(Pid) of
        {ok, #{traps_exits := true}} ->
            [];
        {ok, #{links := Links, traps_exits := false}} ->
            Links;
        _ ->
            []
    end.

-spec monitor_killers(Pid :: pid()) -> [pid()].
monitor_killers(Pid) ->
    case monitor_info(Pid) of
        {ok, #{monitors := Monitors}} ->
            Monitors;
        _ ->
            []
    end.

-spec pids(Children :: [{term(), pid()}]) -> [pid()].
pids(Children) ->
    [Pid || {_Id, Pid} <- Children].

-spec child_sublist(Children :: [{term(), pid()}], Start :: pos_integer(), Len :: pos_integer()) ->
    [{term(), pid()}].
child_sublist(Children, Start, Len) ->
    child_sublist_skip(Children, Start - 1, Len).

child_sublist_skip(Children, 0, Len) ->
    child_prefix(Children, Len);
child_sublist_skip([], _, _) ->
    [];
child_sublist_skip([_ | T], N, Len) ->
    child_sublist_skip(T, N - 1, Len).

-spec child_prefix(Children :: [{term(), pid()}], Len :: non_neg_integer()) -> [{term(), pid()}].
child_prefix(_Children, 0) -> [];
child_prefix([], _) -> [];
child_prefix([H | T], N) -> [H | child_prefix(T, N - 1)].

-spec child_suffix(Children :: [{term(), pid()}], Skip :: non_neg_integer()) -> [{term(), pid()}].
child_suffix(Children, 0) -> Children;
child_suffix([], _) -> [];
child_suffix([_ | T], N) -> child_suffix(T, N - 1).