amplified_nested

defmodule Amplified.Nested do
  @moduledoc ~S'''
  Utilities for working with hierarchical (tree-structured) data.

  Many applications have self-referential schemas — categories that contain
  subcategories, comments with replies, organisational units, tag taxonomies,
  and so on. These are typically stored as flat rows with a `:parent_id`
  foreign key and a `:children` association. `Amplified.Nested` provides a
  small, composable toolkit for the operations you almost always need on
  these structures: nesting a flat list into a tree, flattening a tree back
  into a list, walking ancestors and descendants, and filtering while
  preserving the ancestor chain.

  ## Installation

  Add `amplified_nested` to your list of dependencies in `mix.exs`:

      def deps do
        [
          {:amplified_nested, "~> 0.1.0"}
        ]
      end

  ## Configuration

  Configure the Ecto repo used for preloading (required only if you call
  `descendant_ancestry/3` on structs whose `:children` association has not
  been loaded):

      # config/config.exs
      config :amplified_nested, repo: MyApp.Repo

  If you never use `descendant_ancestry/3` with unloaded associations, no
  configuration is needed.

  ## Setup

  Schema modules opt in with `use Amplified.Nested`:

      defmodule MyApp.Tags.Tag do
        use Ecto.Schema
        use Amplified.Nested

        schema "tags" do
          field :name, :string
          belongs_to :parent, __MODULE__
          has_many :children, __MODULE__, foreign_key: :parent_id
        end
      end

  This generates delegating functions on the schema module itself so you
  can call `Tag.nest(tags)` or `Tag.flatten(tag)` directly, with the
  `:id`, `:parent_id`, and `:children` field names baked in as defaults.

  > #### Always call functions on your schema module {: .info}
  >
  > All functions should be called through the schema module that
  > `use`s `Amplified.Nested` — for example `Tag.nest(tags)`, not
  > `Amplified.Nested.nest(tags)`. The schema delegates handle field
  > name resolution automatically; calling `Amplified.Nested` directly
  > bypasses this and forces you to pass field names explicitly every
  > time.

  ### Custom field names

  If your schema uses non-standard field names, pass them as options:

      use Amplified.Nested, parent: :parent_tag_id, child: :sub_tags, id: :tag_id

  ## Wrapper transparency

  Every function handles `{:ok, ...}`, `{:error, ...}`,
  `Phoenix.LiveView.AsyncResult`, and `Stream` wrappers transparently —
  you can pass them directly from an `assign_async` result or a Repo call
  without unwrapping first:

      result = Tag |> Repo.all() |> then(&{:ok, &1})
      {:ok, nested} = Tag.nest(result)

      # AsyncResult from assign_async
      Tag.nest(socket.assigns.tags)
      #=> %AsyncResult{ok?: true, result: [nested tags...]}

  Error tuples pass through unchanged:

      Tag.nest({:error, :not_found})
      #=> {:error, :not_found}

  ## Functions at a glance

    * `nest/4` — converts a flat list into a parent → children tree
    * `flatten/2` — walks a tree depth-first, returning a flat list
    * `descendant_ancestry/3` — returns each descendant with its full
      ancestry chain (useful for breadcrumbs)
    * `ancestors/5` — finds the ancestry chain for a single entity from
      a flat list
    * `descendants/5` — finds all descendants of an entity from a flat list
    * `filter/5` — filters entities by a predicate, preserving ancestors
      of every match
  '''

  alias Ecto.Association.NotLoaded
  alias Phoenix.LiveView.AsyncResult

  @typep async_result :: %AsyncResult{
           ok?: boolean(),
           loading: term(),
           failed: term(),
           result: term()
         }
  @typep struct_or_structs :: struct() | list(struct())

  @callback nest(
              entities ::
                list(struct()) | async_result() | {:ok, list(struct())} | {:error, term()},
              id :: atom(),
              parent :: atom(),
              child :: atom()
            ) :: list(struct()) | async_result() | {:ok, list(struct())} | {:error, term()}
  @callback flatten(
              entities :: struct_or_structs(),
              child :: atom()
            ) :: list(struct()) | async_result()
  @callback descendant_ancestry(
              entity :: struct() | async_result(),
              parents :: list(struct()),
              child :: atom()
            ) :: list(list(struct())) | async_result()
  @callback ancestors(
              entities :: list(struct()) | async_result(),
              entity :: struct(),
              id :: atom(),
              parent :: atom(),
              acc :: list(struct())
            ) :: list(struct())
  @callback descendants(
              entities :: list(struct()) | async_result(),
              entity :: struct(),
              id :: atom(),
              parent :: atom(),
              acc :: list(struct())
            ) :: list(struct())
  @callback filter(
              entities :: list(struct()) | async_result(),
              filter :: fun(),
              id :: atom(),
              child :: atom(),
              parent :: atom()
            ) :: list(struct())

  @doc """
  Returns the configured Ecto repo module.

  Looked up at runtime via `Application.fetch_env!/2`. Only required if
  you call `descendant_ancestry/3` on structs with unloaded associations.

  Raises `ArgumentError` if `:repo` is not configured for `:amplified_nested`.
  """
  def repo, do: Application.fetch_env!(:amplified_nested, :repo)

  @doc ~S'''
  Injects `Amplified.Nested` behaviour and delegating functions into the
  calling module.

  The generated functions delegate to `Amplified.Nested` with your schema's
  field names baked in as defaults, so you can call `Tag.nest(tags)` instead
  of `Amplified.Nested.nest(tags, :id, :parent_id, :children)`.

  ## Options

    * `:id` — the primary key field name. Defaults to `:id`.
    * `:parent` — the foreign key pointing to the parent. Defaults to
      `:parent_id`.
    * `:child` — the association name for children. Defaults to `:children`.

  ## Examples

  With defaults:

      defmodule MyApp.Tags.Tag do
        use Ecto.Schema
        use Amplified.Nested

        schema "tags" do
          field :name, :string
          belongs_to :parent, __MODULE__
          has_many :children, __MODULE__, foreign_key: :parent_id
        end
      end

      # Now you can call:
      Tag.nest(tags)
      Tag.flatten(tag)
      Tag.ancestors(tags, grandchild)
      Tag.descendants(tags, root)

  With custom field names:

      defmodule MyApp.Org.Unit do
        use Ecto.Schema
        use Amplified.Nested, parent: :reports_to_id, child: :direct_reports

        schema "org_units" do
          field :name, :string
          belongs_to :reports_to, __MODULE__
          has_many :direct_reports, __MODULE__, foreign_key: :reports_to_id
        end
      end

  All generated functions are overridable, so you can provide custom
  implementations for specific schemas when needed.
  '''
  defmacro __using__(opts) do
    parent = Keyword.get(opts, :parent, :parent_id)
    child = Keyword.get(opts, :child, :children)
    id = Keyword.get(opts, :id, :id)
    alias __MODULE__

    quote do
      @behaviour Nested
      @id unquote(id)
      @parent unquote(parent)
      @child unquote(child)

      @impl true
      defdelegate nest(entities, id \\ @id, parent \\ @parent, child \\ @child), to: Nested

      @impl true
      defdelegate flatten(entity, child \\ @child), to: Nested

      @impl true
      defdelegate descendant_ancestry(entity, parents \\ [], child \\ @child), to: Nested

      @impl true
      defdelegate ancestors(entities, entity, id \\ @id, parent \\ @parent, acc \\ []), to: Nested

      @impl true
      defdelegate descendants(entities, entity, id \\ @id, parent \\ @parent, acc \\ []),
        to: Nested

      @impl true
      defdelegate filter(entities, filter, id \\ @id, child \\ @child, parent \\ @parent),
        to: Nested

      defoverridable Nested
    end
  end

  @doc ~S'''
  Converts a flat list of structs into a parent → children tree.

  Takes a list of structs that each have an `:id` and `:parent_id` field
  (or equivalents), groups them by parent, and recursively populates the
  `:children` field to build a nested tree. Root nodes are those with a
  `nil` parent.

  Returns only the root nodes, with all descendants nested under them.

  ## Parameters

    * `entities` — a flat list of structs, or an `{:ok, list}` /
      `{:error, reason}` / `AsyncResult` / `Stream` wrapper.
    * `id` — the field name used as the primary key. Defaults to `:id`.
    * `parent` — the field name used as the parent foreign key. Defaults
      to `:parent_id`.
    * `child` — the field name for the children association. Defaults to
      `:children`.

  ## Examples

  Basic nesting (called via schema delegate):

      tags = [
        %Tag{id: 1, parent_id: nil, name: "Elixir"},
        %Tag{id: 2, parent_id: nil, name: "Rust"},
        %Tag{id: 3, parent_id: 1,   name: "Phoenix"},
        %Tag{id: 4, parent_id: 3,   name: "LiveView"}
      ]

      Tag.nest(tags)
      #=> [
      #     %Tag{id: 1, name: "Elixir", children: [
      #       %Tag{id: 3, name: "Phoenix", children: [
      #         %Tag{id: 4, name: "LiveView", children: []}
      #       ]}
      #     ]},
      #     %Tag{id: 2, name: "Rust", children: []}
      #   ]

  Pipeline with `{:ok, _}` wrapper:

      Tag.nest({:ok, tags})
      #=> {:ok, [%Tag{...nested...}]}

  Error tuples pass through:

      Tag.nest({:error, :not_found})
      #=> {:error, :not_found}

  Single-element lists are returned as-is (no nesting needed):

      Tag.nest([%Tag{id: 1}])
      #=> [%Tag{id: 1}]
  '''
  @spec nest(
          entities :: list(struct()) | async_result() | {:ok, list(struct())} | {:error, term()},
          id :: atom(),
          parent :: atom(),
          child :: atom()
        ) :: list(struct()) | async_result() | {:ok, list(struct())} | {:error, term()}
  def nest(entities, id \\ :id, parent \\ :parent_id, child \\ :children)
  def nest({:ok, entities}, id, parent, child), do: {:ok, nest(entities, id, parent, child)}
  def nest({:error, error}, _id, _parent, _child), do: {:error, error}

  def nest(%AsyncResult{ok?: true, result: entities}, id, parent, child),
    do: entities |> nest(id, parent, child) |> AsyncResult.ok()

  def nest(%Stream{} = stream, id, parent, child),
    do: stream |> Enum.to_list() |> nest(id, parent, child)

  def nest([_] = entities, _, _, _), do: entities

  def nest([entity | _] = entities, id, parent, child)
      when is_map_key(entity, parent),
      do: entities |> Enum.group_by(&Map.get(&1, parent)) |> nest(nil, id, parent, child)

  def nest(entities, _id, _parent, _child), do: entities

  defp nest(entities, parent_id, id, parent, child) do
    entities
    |> Map.get(parent_id, [])
    |> Enum.map(&Map.put(&1, child, nest(entities, Map.get(&1, id), id, parent, child)))
  end

  @doc ~S'''
  Flattens a nested tree structure into a depth-first ordered list.

  Takes a struct (or list of structs) with a populated `:children` field
  and returns every node in the tree as a flat list, traversed depth-first.
  This is the inverse of `nest/4`.

  Handles `Ecto.Association.NotLoaded` children gracefully — those nodes
  are treated as leaf nodes.

  ## Parameters

    * `entity` — a struct, list of structs, `AsyncResult`, or `Stream`.
    * `child` — the field name for the children association. Defaults to
      `:children`. Pass `nil` to return the entity unchanged.

  ## Examples

  Flatten a single nested struct:

      tag = %Tag{
        id: 1, name: "Elixir",
        children: [
          %Tag{id: 2, name: "Phoenix", children: [
            %Tag{id: 3, name: "LiveView", children: []}
          ]},
          %Tag{id: 4, name: "Ecto", children: []}
        ]
      }

      Tag.flatten(tag)
      #=> [
      #     %Tag{id: 1, name: "Elixir", ...},
      #     %Tag{id: 2, name: "Phoenix", ...},
      #     %Tag{id: 3, name: "LiveView", ...},
      #     %Tag{id: 4, name: "Ecto", ...}
      #   ]

  Flatten a list of root nodes:

      Tag.flatten([root1, root2])
      #=> [root1, root1_child, root1_grandchild, root2, root2_child, ...]
  '''
  @spec flatten(
          entities :: struct_or_structs(),
          child :: atom()
        ) :: list(struct()) | async_result()
  def flatten(entity, child \\ :children)
  def flatten(entity, nil), do: entity

  def flatten(%AsyncResult{ok?: true, result: entity}, child),
    do: entity |> flatten(child) |> AsyncResult.ok()

  def flatten(%Stream{} = stream, child), do: stream |> Enum.to_list() |> flatten(child)
  def flatten(%{} = entity, child), do: flatten([entity], child)

  def flatten(entities, child) when is_list(entities) do
    Enum.flat_map(entities, fn
      %{^child => []} = entity -> [entity]
      %{^child => %NotLoaded{}} = entity -> [entity]
      %{^child => children} = entity -> [entity | flatten(children, child)]
      entity -> [entity]
    end)
  end

  def flatten(_, _), do: []

  @doc ~S'''
  Returns each descendant of the given entity paired with its full ancestry.

  Walks the tree depth-first and produces a list of lists, where each inner
  list is a "breadcrumb trail" from the root down to that node. The first
  entry is the root itself (ancestry of length 1), the second is
  `[root, child]`, and so on.

  If a node's `:children` association is `%Ecto.Association.NotLoaded{}`,
  the configured `repo` is used to preload it. If you don't need this
  behaviour, ensure children are preloaded before calling.

  ## Parameters

    * `entity` — a single struct (the root), or an `AsyncResult` wrapper.
    * `parents` — the accumulated parent chain. You normally don't pass
      this — it defaults to `[]` and is built up during recursion.
    * `child` — the children field name. Defaults to `:children`.

  ## Examples

      tag = %Tag{
        id: 1, name: "Elixir",
        children: [
          %Tag{id: 2, name: "Phoenix", children: [
            %Tag{id: 3, name: "LiveView", children: []}
          ]},
          %Tag{id: 4, name: "Ecto", children: []}
        ]
      }

      Tag.descendant_ancestry(tag)
      #=> [
      #     [%Tag{name: "Elixir"}],
      #     [%Tag{name: "Elixir"}, %Tag{name: "Phoenix"}],
      #     [%Tag{name: "Elixir"}, %Tag{name: "Phoenix"}, %Tag{name: "LiveView"}],
      #     [%Tag{name: "Elixir"}, %Tag{name: "Ecto"}]
      #   ]

  This is particularly useful for rendering breadcrumb navigation or
  indented tree views where each row needs to know its full path from
  the root.
  '''
  @spec descendant_ancestry(
          entity :: struct() | async_result(),
          parents :: list(struct()) | async_result(),
          child :: atom()
        ) :: list(list(struct()))
  def descendant_ancestry(entity, parents \\ [], child \\ :children)

  def descendant_ancestry(%AsyncResult{ok?: true, result: entity}, parents, child),
    do: entity |> descendant_ancestry(parents, child) |> AsyncResult.ok()

  def descendant_ancestry(entity, parents, child) when is_map(entity) do
    children =
      case Map.get(entity, child) do
        %NotLoaded{} -> entity |> repo().preload(child) |> Map.get(child)
        children -> children
      end

    parents = parents ++ [entity]
    [parents | Enum.flat_map(children || [], &descendant_ancestry(&1, parents, child))]
  end

  @doc ~S'''
  Returns the ancestry chain for a given entity from a flat list.

  Walks up the parent chain from the given entity, collecting each
  ancestor, and returns the full path from the oldest ancestor down to the
  entity itself. The list is ordered root-first.

  This operates on a **flat** list (not a nested tree) — it looks up
  parents by matching `:id` and `:parent_id` fields.

  ## Parameters

    * `entities` — the flat list of all entities (or an `AsyncResult`).
    * `entity` — the entity whose ancestry you want.
    * `id` — the primary key field. Defaults to `:id`.
    * `parent` — the parent foreign key field. Defaults to `:parent_id`.
    * `acc` — accumulator, defaults to `[]`. Typically not passed by the
      caller.

  ## Examples

      tags = [
        %Tag{id: 1, parent_id: nil, name: "Elixir"},
        %Tag{id: 2, parent_id: 1,   name: "Phoenix"},
        %Tag{id: 3, parent_id: 2,   name: "LiveView"}
      ]

      liveview = Enum.find(tags, &(&1.id == 3))

      Tag.ancestors(tags, liveview)
      #=> [
      #     %Tag{id: 1, name: "Elixir"},
      #     %Tag{id: 2, name: "Phoenix"},
      #     %Tag{id: 3, name: "LiveView"}
      #   ]

  A root entity (no parent) returns just itself:

      root = Enum.find(tags, &(&1.id == 1))
      Tag.ancestors(tags, root)
      #=> [%Tag{id: 1, name: "Elixir"}]
  '''
  @spec ancestors(
          entities :: list(struct()) | async_result(),
          entity :: struct(),
          id :: atom(),
          parent :: atom(),
          acc :: list(struct())
        ) :: list(struct())
  def ancestors(entities, entity, id \\ :id, parent \\ :parent_id, acc \\ [])

  def ancestors(%AsyncResult{ok?: true, result: entities}, entity, id, parent, acc),
    do: ancestors(entities, entity, id, parent, acc)

  def ancestors(entities, entity, id, parent, acc) when is_list(entities) do
    entities
    |> Map.new(&{Map.get(&1, id), &1})
    |> ancestors(entity, id, parent, acc)
    |> Enum.filter(& &1)
  end

  def ancestors(entities, entity, id, parent, acc) when is_map(entities) do
    parent_id = Map.get(entity, parent)

    if not is_map_key(entities, :result) and not is_nil(parent_id),
      do: ancestors(entities, Map.get(entities, parent_id), id, parent, [entity | acc]),
      else: [entity | acc]
  end

  @doc ~S'''
  Returns all descendants of the given entity from a flat list.

  Walks down the tree by finding entities whose `:parent_id` matches the
  given entity's `:id`, recursing into each child. The given entity itself
  is included in the result.

  This operates on a **flat** list (not a nested tree).

  ## Parameters

    * `entities` — the flat list of all entities (or an `AsyncResult`).
    * `entity` — the entity whose descendants you want.
    * `id` — the primary key field. Defaults to `:id`.
    * `parent` — the parent foreign key field. Defaults to `:parent_id`.
    * `acc` — accumulator, defaults to `[]`. Typically not passed by the
      caller.

  ## Examples

      tags = [
        %Tag{id: 1, parent_id: nil, name: "Elixir"},
        %Tag{id: 2, parent_id: 1,   name: "Phoenix"},
        %Tag{id: 3, parent_id: 2,   name: "LiveView"},
        %Tag{id: 4, parent_id: nil, name: "Rust"}
      ]

      elixir = Enum.find(tags, &(&1.id == 1))

      Tag.descendants(tags, elixir)
      #=> [
      #     %Tag{id: 2, name: "Phoenix"},
      #     %Tag{id: 3, name: "LiveView"},
      #     %Tag{id: 1, name: "Elixir"}
      #   ]

  Note that the entity itself appears at the end of the list. Entities
  from other branches (like "Rust" above) are excluded.

  A leaf entity with no children returns just itself:

      liveview = Enum.find(tags, &(&1.id == 3))
      Tag.descendants(tags, liveview)
      #=> [%Tag{id: 3, name: "LiveView"}]
  '''
  @spec descendants(
          entities :: list(struct()) | async_result(),
          entity :: struct(),
          id :: atom(),
          parent :: atom(),
          acc :: list(struct())
        ) :: list(struct())
  def descendants(entities, entity, id \\ :id, parent \\ :parent_id, acc \\ [])

  def descendants(%AsyncResult{ok?: true, result: entities}, entity, id, parent, acc)
      when is_list(entities),
      do: descendants(entities, entity, id, parent, acc)

  def descendants(entities, entity, id, parent, acc) do
    entity_id = Map.get(entity, id)

    entities
    |> Enum.filter(&(Map.get(&1, parent) == entity_id))
    |> case do
      [] -> acc
      children -> Enum.map(children, &descendants(entities, &1, id, parent, acc))
    end
    |> Enum.concat([entity])
    |> List.flatten()
  end

  @doc ~S'''
  Filters entities by a predicate, preserving the ancestor chain of every
  match.

  First flattens the input (in case it's a nested tree), then applies the
  filter function. For every entity that matches, its full ancestry is
  included in the result. Duplicates (shared ancestors) are removed by
  `:id`.

  This is useful for search — when a user searches for "LiveView" in a
  tag tree, you want to show "LiveView" but also its parent "Phoenix" and
  grandparent "Elixir" so the user can see the context.

  ## Parameters

    * `entities` — a nested tree or flat list (or `AsyncResult`).
    * `filter` — a function that returns truthy for entities to keep.
    * `id` — the primary key field. Defaults to `:id`.
    * `child` — the children field name. Defaults to `:children`.
    * `parent` — the parent foreign key field. Defaults to `:parent_id`.

  ## Examples

      tags = [
        %Tag{id: 1, name: "Elixir", parent_id: nil, children: [
          %Tag{id: 2, name: "Phoenix", parent_id: 1, children: [
            %Tag{id: 3, name: "LiveView", parent_id: 2, children: []}
          ]}
        ]},
        %Tag{id: 4, name: "Rust", parent_id: nil, children: []}
      ]

      Tag.filter(tags, &(&1.name == "LiveView"))
      #=> [
      #     %Tag{id: 1, name: "Elixir"},
      #     %Tag{id: 2, name: "Phoenix"},
      #     %Tag{id: 3, name: "LiveView"}
      #   ]

  Note that "Rust" is excluded because it's not an ancestor of any match.
  '''
  @spec filter(
          entities :: list(struct()) | async_result(),
          filter :: fun(),
          id :: atom(),
          child :: atom(),
          parent :: atom()
        ) :: list(struct())
  def filter(entities, filter, id \\ :id, child \\ :children, parent \\ :parent_id)

  def filter(%AsyncResult{ok?: true, result: entities}, filter, id, child, parent),
    do: entities |> filter(filter, id, child, parent) |> AsyncResult.ok()

  def filter(entities, filter, id, child, parent) do
    entities = flatten(entities, child)

    entities
    |> Stream.filter(filter)
    |> Stream.flat_map(&ancestors(entities, &1, id, parent, []))
    |> Enum.uniq_by(&Map.get(&1, id))
  end
end