defmodule Firebreak.Forest do
  @moduledoc """
  Builds the supervision forest from parsed modules: who supervises whom, which
  supervisors are roots, and the full set of descendant modules in each
  supervisor's subtree.

  The subtree is what Tier-2 blast-radius analysis is computed against — a
  coupling edge that crosses *out* of a subtree is the interesting one.
  """

  alias Firebreak.ModuleInfo

  @type t :: %{
          supervisors: [module()],
          roots: [module()],
          child_map: %{optional(module()) => [module()]},
          parents: %{optional(module()) => [module()]},
          subtree: %{optional(module()) => MapSet.t()},
          child_specs: %{optional(module()) => Firebreak.Child.t()},
          strategies: %{optional(module()) => atom()}
        }

  @spec build([ModuleInfo.t()]) :: t()
  def build(modules) do
    by_name = Map.new(modules, &{&1.name, &1})
    supervisors = for m <- modules, ModuleInfo.supervisor?(m), do: m.name

    child_map =
      Map.new(supervisors, fn sup ->
        info = Map.fetch!(by_name, sup)

        # A DynamicSupervisor's children are added at runtime, never in init/1, so
        # we fold in the best-effort `dynamic_children` recovered from start_child
        # call sites. They become subtree members like any other child, so a
        # caller crossing into them is caught by blast-radius analysis.
        children =
          (child_modules(info.children) ++ child_modules(info.dynamic_children))
          |> Enum.uniq()

        {sup, children}
      end)

    parents = invert(child_map)

    roots =
      Enum.filter(supervisors, fn sup ->
        app?(by_name[sup]) or not parented?(sup, parents)
      end)

    subtree = Map.new(supervisors, fn sup -> {sup, descendants(sup, child_map)} end)

    %{
      supervisors: supervisors,
      roots: roots,
      child_map: child_map,
      parents: parents,
      subtree: subtree,
      child_specs: child_specs(supervisors, by_name),
      strategies: Map.new(supervisors, fn sup -> {sup, by_name[sup].strategy} end)
    }
  end

  defp child_modules(children) do
    children |> Enum.map(& &1.module) |> Enum.reject(&is_nil/1)
  end

  # Index each child module to its spec, so blast-radius analysis can read the
  # child's restart type (`:permanent` / `:transient` / `:temporary`). A module
  # supervised in more than one place keeps the first spec that names a restart,
  # since that is the entry that changes the failure story.
  defp child_specs(supervisors, by_name) do
    Enum.reduce(supervisors, %{}, fn sup, acc ->
      info = Map.fetch!(by_name, sup)

      (info.children ++ info.dynamic_children)
      |> Enum.reduce(acc, fn
        %{module: nil}, acc2 -> acc2
        %{module: mod} = child, acc2 -> Map.update(acc2, mod, child, &prefer_restart(&1, child))
      end)
    end)
  end

  defp prefer_restart(%{restart: nil}, %{restart: r} = newer) when not is_nil(r), do: newer
  defp prefer_restart(existing, _newer), do: existing

  defp app?(%ModuleInfo{kind: :application}), do: true
  defp app?(_), do: false

  # A supervisor is "parented" if some *other* supervisor lists it as a child.
  defp parented?(sup, parents) do
    parents |> Map.get(sup, []) |> Enum.any?(&(&1 != sup))
  end

  defp invert(child_map) do
    Enum.reduce(child_map, %{}, fn {sup, children}, acc ->
      Enum.reduce(children, acc, fn child, acc2 ->
        Map.update(acc2, child, [sup], &[sup | &1])
      end)
    end)
  end

  # All modules reachable downward from `sup` (including sup itself), following
  # child links through sub-supervisors. Cycle-guarded.
  defp descendants(sup, child_map) do
    walk([sup], child_map, MapSet.new())
  end

  defp walk([], _child_map, seen), do: seen

  defp walk([node | rest], child_map, seen) do
    if MapSet.member?(seen, node) do
      walk(rest, child_map, seen)
    else
      seen = MapSet.put(seen, node)
      children = Map.get(child_map, node, [])
      walk(children ++ rest, child_map, seen)
    end
  end
end