defmodule Ash.Actions.Helpers do
  @moduledoc false
  require Logger
  require Ash.Flags
  alias Ash.Actions.BulkManualActionHelpers

  @keep_read_action_loads_when_loading? Application.compile_env(
                                          :ash,
                                          :keep_read_action_loads_when_loading?,
                                          true
                                        )

  @spec keep_read_action_loads_when_loading? :: boolean()
  def keep_read_action_loads_when_loading?, do: @keep_read_action_loads_when_loading?

  def split_and_run_simple(batch, action, opts, changes, all_changes, context_key, callback) do
    {batch, must_be_simple} =
      Enum.reduce(batch, {[], []}, fn changeset, {batch, must_be_simple} ->
        if changeset.around_transaction in [[], nil] and changeset.after_transaction in [[], nil] and
             changeset.around_action in [[], nil] do
          changeset =
            if changeset.valid? do
              Ash.Changeset.run_before_transaction_hooks(changeset)
            else
              changeset
            end

          {[changeset | batch], must_be_simple}
        else
          {batch, [%{changeset | __validated_for_action__: action.name} | must_be_simple]}
        end
      end)

    batch = batch |> Enum.reverse()
    must_be_simple = must_be_simple |> Enum.reverse()

    context =
      case {batch, must_be_simple} do
        {[cs | _], _} -> cs.context
        {_, [cs | _]} -> cs.context
        {_, _} -> %{}
      end

    context =
      struct(
        Ash.Resource.Change.Context,
        %{
          bulk?: true,
          source_context: context,
          actor: opts[:actor],
          tenant: opts[:tenant],
          tracer: opts[:tracer],
          authorize?: opts[:authorize?]
        }
      )

    must_be_simple_results =
      Enum.flat_map(must_be_simple, fn changeset ->
        changeset =
          all_changes
          |> Enum.flat_map(fn
            {%{change: {mod, change_opts}} = change, change_index} ->
              index =
                BulkManualActionHelpers.get_changeset_index(changeset, context_key)

              applicable = changes[change_index]

              if applicable == :all || (applicable && index in applicable) do
                change_opts =
                  Ash.Actions.Helpers.templated_opts(
                    change_opts,
                    opts[:actor],
                    changeset.to_tenant,
                    changeset.arguments,
                    changeset.context,
                    changeset
                  )

                if mod.batch_callbacks?([changeset], change_opts, context) do
                  [%{change | change: {mod, change_opts}}]
                else
                  []
                end
              else
                []
              end

            _ ->
              []
          end)
          |> Enum.reduce(changeset, fn %{change: {mod, change_opts}}, changeset ->
            changeset =
              if mod.has_after_batch?() do
                Ash.Changeset.after_action(changeset, fn changeset, result ->
                  case mod.after_batch([{changeset, result}], change_opts, context) do
                    :ok ->
                      {:ok, result}

                    enumerable ->
                      Enum.reduce_while(
                        enumerable,
                        {{:ok, result}, []},
                        fn
                          %Ash.Notifier.Notification{} = notification, {res, notifications} ->
                            {:cont, {res, [notification | notifications]}}

                          {:error, error}, {_res, notifications} ->
                            {:halt, {{:error, error}, notifications}}

                          {:ok, result}, {_res, notifications} ->
                            {:cont, {{:ok, result}, notifications}}
                        end
                      )
                  end
                  |> case do
                    {{:ok, res}, notifications} -> {:ok, res, notifications}
                    {other, _} -> other
                  end
                end)
              else
                changeset
              end

            if mod.has_before_batch?() do
              Ash.Changeset.before_action(changeset, fn changeset ->
                mod.before_batch([changeset], change_opts, context)
                |> Enum.reduce(
                  {changeset, []},
                  fn
                    %Ash.Notifier.Notification{} = notification, {changeset, notifications} ->
                      {changeset, [notification | notifications]}

                    changeset, {_, notifications} ->
                      {changeset, notifications}
                  end
                )
              end)
            else
              changeset
            end
          end)

        callback.(changeset)
      end)

    {batch, must_be_simple_results}
  end

  def rollback_if_in_transaction(
        {:error, %Ash.Error.Changes.StaleRecord{} = error},
        _resource,
        _changeset
      ) do
    {:error, error}
  end

  def rollback_if_in_transaction({:error, error}, resource, changeset) do
    error = Ash.Error.to_ash_error(error)

    if Ash.DataLayer.in_transaction?(resource) do
      case changeset do
        %Ash.Changeset{} = changeset ->
          Ash.DataLayer.rollback(resource, Ash.Changeset.add_error(changeset, error))

        %Ash.Query{} = query ->
          Ash.DataLayer.rollback(resource, Ash.Query.add_error(query, error))

        %Ash.ActionInput{} = action_input ->
          Ash.DataLayer.rollback(resource, Ash.ActionInput.add_error(action_input, error))

        _ ->
          Ash.DataLayer.rollback(resource, Ash.Error.to_error_class(error))
      end
    else
      {:error, error}
    end
  end

  def rollback_if_in_transaction({:error, :no_rollback, error}, _, _changeset) do
    {:error, error}
  end

  def rollback_if_in_transaction(success, _, _), do: success

  def validate_calculation_load!(%Ash.Query{}, module) do
    raise """
    `#{inspect(module)}.load/3` returned a query.

    Returning a query from the `load/3` callback of a calculation is now deprecated.
    Instead, return the load statement itself, i.e instead of `Ash.Query.load(query, [...])`,
    just return `[...]`. This is so that Ash can examine the requirements of just this single
    calculation to ensure that all required values are present
    """
  end

  def validate_calculation_load!(other, _), do: List.wrap(other)

  defp set_skip_unknown_opts(opts, %{action: %{skip_unknown_inputs: skip_unknown_inputs}}) do
    Keyword.update(
      opts,
      :skip_unknown_inputs,
      skip_unknown_inputs,
      &Enum.concat(List.wrap(&1), skip_unknown_inputs)
    )
  end

  defp set_skip_unknown_opts(opts, _query_or_changeset) do
    opts
  end

  def maybe_embedded_domain(resource) do
    if Ash.Resource.Info.embedded?(resource) do
      Ash.EmbeddableType.ShadowDomain
    end
  end

  def apply_scope_to_opts(opts) do
    if scope = opts[:scope] do
      actor = Ash.Scope.ToOpts.get_actor(scope)
      tenant = Ash.Scope.ToOpts.get_tenant(scope)
      context = Ash.Scope.ToOpts.get_context(scope)
      authorize? = Ash.Scope.ToOpts.get_authorize?(scope)

      opts
      |> set_when_ok(:actor, actor, fn l, _r -> l end)
      |> set_when_ok(:tenant, tenant, fn l, _r -> l end)
      |> set_when_ok(:authorize?, authorize?, fn l, _r -> l end)
      |> set_when_ok(
        :context,
        context,
        &Ash.Helpers.deep_merge_maps(&2, &1)
      )
      |> Keyword.delete(:scope)
    else
      opts
    end
  end

  def set_context_and_get_opts(domain, query_or_changeset, opts) do
    opts = apply_scope_to_opts(opts)

    opts = set_skip_unknown_opts(opts, query_or_changeset)
    query_or_changeset = Ash.Subject.set_context(query_or_changeset, opts[:context] || %{})

    domain =
      Ash.Resource.Info.domain(query_or_changeset.resource) || opts[:domain] || domain ||
        query_or_changeset.domain || maybe_embedded_domain(query_or_changeset.resource)

    opts =
      case query_or_changeset.context do
        %{
          private: %{
            actor: actor
          }
        } ->
          Keyword.put_new(opts, :actor, actor)

        _ ->
          opts
      end

    opts =
      if tenant = query_or_changeset.tenant do
        Keyword.put_new(opts, :tenant, tenant)
      else
        opts
      end

    opts =
      case query_or_changeset.context do
        %{
          private: %{
            authorize?: authorize?
          }
        }
        when not is_nil(authorize?) ->
          Keyword.put_new(opts, :authorize?, authorize?)

        _ ->
          opts
      end

    opts =
      case query_or_changeset.context do
        %{
          private: %{
            tracer: tracer
          }
        } ->
          do_add_tracer(opts, tracer)

        _ ->
          opts
      end

    opts = set_opts(opts, domain, query_or_changeset)

    query_or_changeset = add_context(query_or_changeset, opts)

    query_or_changeset = %{query_or_changeset | domain: domain}

    {query_or_changeset, opts}
  end

  @doc false
  def set_when_ok(opts, key, value, merger \\ fn _l, r -> r end)

  def set_when_ok(opts, key, {:ok, value}, merger) do
    Keyword.update(opts, key, value, &merger.(&1, value))
  end

  def set_when_ok(opts, _, _, _), do: opts

  def set_opts(opts, domain, query_or_changeset \\ nil) do
    opts
    |> add_actor(query_or_changeset, domain)
    |> add_authorize?(query_or_changeset, domain)
    |> add_tracer()
  end

  def add_context(query_or_changeset, opts) do
    private_context = Map.new(Keyword.take(opts, [:actor, :authorize?, :tracer]))

    case query_or_changeset do
      %Ash.ActionInput{} ->
        query_or_changeset
        |> Ash.ActionInput.set_context(%{private: private_context})
        |> Ash.ActionInput.set_tenant(query_or_changeset.tenant || opts[:tenant])

      %Ash.Query{} ->
        query_or_changeset
        |> Ash.Query.set_context(%{private: private_context})
        |> Ash.Query.set_tenant(query_or_changeset.tenant || opts[:tenant])

      %Ash.Changeset{} ->
        query_or_changeset
        |> Ash.Changeset.set_context(%{
          private: private_context
        })
        |> Ash.Changeset.set_tenant(query_or_changeset.tenant || opts[:tenant])
    end
  end

  defp add_actor(opts, query_or_changeset, domain) do
    if !domain do
      raise Ash.Error.Framework.AssumptionFailed,
        message: "Could not determine domain for action."
    end

    if !skip_requiring_actor?(query_or_changeset) && !internal?(query_or_changeset) &&
         !Keyword.has_key?(opts, :actor) &&
         Ash.Domain.Info.require_actor?(domain) do
      raise Ash.Error.to_error_class(
              Ash.Error.Forbidden.DomainRequiresActor.exception(domain: domain)
            )
    end

    opts
  end

  defp internal?(%{context: %{private: %{internal?: true}}}), do: true
  defp internal?(_), do: false

  defp skip_requiring_actor?(%{context: %{private: %{require_actor?: false}}}), do: true
  defp skip_requiring_actor?(_), do: false

  defp add_authorize?(opts, query_or_changeset, domain) do
    if !domain do
      raise Ash.Error.Framework.AssumptionFailed,
        message: "Could not determine domain for action."
    end

    case Ash.Domain.Info.authorize(domain) do
      :always ->
        if opts[:authorize?] == false && internal?(query_or_changeset) do
          opts
        else
          if opts[:authorize?] == false do
            raise Ash.Error.Forbidden.DomainRequiresAuthorization, domain: domain
          end

          Keyword.put(opts, :authorize?, true)
        end

      :by_default ->
        Keyword.put_new(opts, :authorize?, true)

      :when_requested ->
        if Keyword.has_key?(opts, :actor) do
          Keyword.put_new(opts, :authorize?, true)
        else
          Keyword.put(opts, :authorize?, opts[:authorize?] || Keyword.has_key?(opts, :actor))
        end
    end
  end

  defp add_tracer(opts) do
    case Application.get_env(:ash, :tracer) do
      nil ->
        opts

      tracer ->
        do_add_tracer(opts, tracer)
    end
  end

  defp do_add_tracer(opts, tracer) do
    tracer = List.wrap(tracer)

    Keyword.update(opts, :tracer, tracer, fn existing_tracer ->
      if is_list(existing_tracer) do
        Enum.uniq(tracer ++ existing_tracer)
      else
        if is_nil(existing_tracer) do
          tracer
        else
          Enum.uniq(tracer ++ existing_tracer)
        end
      end
    end)
  end

  @doc false
  def notify({:ok, record, instructions}, changeset, opts) do
    resource_notification = resource_notification(changeset, record, opts)

    if opts[:return_notifications?] do
      {:ok, record,
       Map.update(
         instructions,
         :notifications,
         [resource_notification],
         &[resource_notification | &1]
       )}
    else
      if Process.get(:ash_started_transaction?) do
        current_notifications = Process.get(:ash_notifications, [])

        Process.put(
          :ash_notifications,
          [resource_notification | current_notifications]
        )
      else
        unsent_notifications = Ash.Notifier.notify([resource_notification])

        warn_missed!(changeset.resource, changeset.action, %{
          resource_notifications: unsent_notifications
        })
      end

      {:ok, record, instructions}
    end
  end

  def notify(other, _changeset, _opts), do: other

  defp resource_notification(changeset, result, opts) do
    # This gives notifications a view of what actually changed
    changeset =
      Enum.reduce(changeset.atomics, changeset, fn {key, value}, changeset ->
        %{changeset | attributes: Map.put(changeset.attributes, key, value)}
      end)

    %Ash.Notifier.Notification{
      resource: changeset.resource,
      domain: changeset.domain,
      actor: changeset.context[:private][:actor],
      action: changeset.action,
      for: Ash.Resource.Info.notifiers(changeset.resource) ++ changeset.action.notifiers,
      data: result,
      changeset: changeset,
      metadata: opts[:notification_metadata] || %{}
    }
  end

  def warn_missed!(resource, action, result) do
    case Map.get(result, :resource_notifications, Map.get(result, :notifications, [])) do
      empty when empty in [nil, []] ->
        :ok

      missed ->
        case Application.get_env(:ash, :missed_notifications, :warn) do
          :ignore ->
            :ok

          :raise ->
            raise """
            Missed #{Enum.count(missed)} notifications in action #{inspect(resource)}.#{action.name}.

            This happens when the resources are in a transaction, and you did not pass
            `return_notifications?: true`. If you are in a changeset hook, you can
            return the notifications. If not, you can send the notifications using
            `Ash.Notifier.notify/1` once your resources are out of a transaction.

            To ignore these in all cases:

            config :ash, :missed_notifications, :ignore

            To turn this into warnings:

            config :ash, :missed_notifications, :warn
            """

          :warn ->
            {:current_stacktrace, stacktrace} = Process.info(self(), :current_stacktrace)

            Logger.warning("""
            Missed #{Enum.count(missed)} notifications in action #{inspect(resource)}.#{action.name}.

            This happens when the resources are in a transaction, and you did not pass
            `return_notifications?: true`. If you are in a changeset hook, you can
            return the notifications. If not, you can send the notifications using
            `Ash.Notifier.notify/1` once your resources are out of a transaction.

            #{Exception.format_stacktrace(stacktrace)}

            While you should likely leave this setting on, you can ignore these or turn them into errors.

            To ignore these in all cases:

            config :ash, :missed_notifications, :ignore

            To turn this into raised errors:

            config :ash, :missed_notifications, :raise
            """)
        end
    end
  end

  def process_errors(changeset, [error]) do
    %{changeset | errors: []}
    |> Ash.Changeset.add_error(error)
    |> Map.get(:errors)
    |> case do
      [error] ->
        error

      errors ->
        errors
    end
  end

  def process_errors(changeset, errors) when is_list(errors) do
    %{changeset | errors: []}
    |> Ash.Changeset.add_error(errors)
    |> Map.get(:errors)
  end

  def process_errors(changeset, error), do: process_errors(changeset, [error])

  def templated_opts({:templated, opts}, _actor, _tenant, _arguments, _context, _changeset),
    do: opts

  def templated_opts(opts, actor, tenant, arguments, context, changeset) do
    Ash.Expr.fill_template(
      opts,
      actor: actor,
      tenant: tenant,
      args: arguments,
      context: context,
      changeset: changeset
    )
  end

  def load_runtime_types({:ok, results}, query, attributes?) do
    load_runtime_types(results, query, attributes?)
  end

  def load_runtime_types({:error, error}, _query, _attributes?) do
    {:error, error}
  end

  def load_runtime_types(results, query, attributes?) when is_list(results) do
    attributes = runtime_attributes(query, attributes?)
    calcs = runtime_calculations(query)

    if Enum.empty?(attributes) && Enum.empty?(calcs) do
      {:ok, results}
    else
      Enum.reduce_while(results, {:ok, []}, fn result, {:ok, results} ->
        case do_load_runtime_types(result, attributes, calcs) do
          {:ok, result} ->
            {:cont, {:ok, [result | results]}}

          {:error, error} ->
            {:halt, {:error, error}}
        end
      end)
      |> case do
        {:ok, results} -> {:ok, Enum.reverse(results)}
        {:error, error} -> {:error, error}
      end
    end
  end

  def load_runtime_types(nil, _, _attributes?), do: {:ok, nil}

  def load_runtime_types(result, query, attributes?) do
    do_load_runtime_types(
      result,
      runtime_attributes(query, attributes?),
      runtime_calculations(query)
    )
  end

  defp runtime_attributes(query, true) do
    case query.select do
      nil ->
        Ash.Resource.Info.attributes(query.resource)

      select ->
        Enum.map(select, &Ash.Resource.Info.attribute(query.resource, &1))
    end
    |> Enum.reject(fn %{type: type, constraints: constraints} ->
      Ash.Type.cast_in_query?(type, constraints)
    end)
  end

  defp runtime_attributes(_, _), do: []

  defp runtime_calculations(query) do
    query.calculations
    |> Kernel.||(%{})
    |> Enum.filter(fn {_name, calc} ->
      calc.type
    end)
    |> Enum.reject(fn {_name, calc} ->
      constraints = Map.get(calc, :constraints, [])

      if function_exported?(Ash.Type, :cast_in_query?, 2) do
        Ash.Type.cast_in_query?(calc.type, constraints)
      else
        Ash.Type.cast_in_query?(calc.type)
      end
    end)
  end

  defp do_load_runtime_types(record, select, calculations) do
    select
    |> Enum.reduce_while({:ok, record}, fn attr, {:ok, record} ->
      case Map.get(record, attr.name) do
        nil ->
          {:cont, {:ok, record}}

        %Ash.NotLoaded{} ->
          {:cont, {:ok, record}}

        %Ash.ForbiddenField{} ->
          {:cont, {:ok, record}}

        value ->
          case Ash.Type.cast_stored(
                 attr.type,
                 value,
                 attr.constraints
               ) do
            {:ok, value} ->
              {:cont, {:ok, Map.put(record, attr.name, value)}}

            :error ->
              {:halt, {:error, message: "is invalid", field: attr.name}}
          end
      end
    end)
    |> case do
      {:ok, record} ->
        Enum.reduce_while(calculations, {:ok, record}, fn {name, calc}, {:ok, record} ->
          case calc.load do
            nil ->
              case Map.get(record.calculations || %{}, calc.name) do
                nil ->
                  {:cont, {:ok, record}}

                value ->
                  case Ash.Type.cast_stored(
                         calc.type,
                         value,
                         Map.get(calc, :constraints, [])
                       ) do
                    {:ok, value} ->
                      {:cont,
                       {:ok, Map.update!(record, :calculations, &Map.put(&1, name, value))}}

                    :error ->
                      {:halt, {:error, message: "is invalid", field: calc.name}}
                  end
              end

            load ->
              case Map.get(record, load) do
                nil ->
                  {:cont, {:ok, record}}

                value ->
                  case Ash.Type.cast_stored(
                         calc.type,
                         value,
                         Map.get(calc, :constraints, [])
                       ) do
                    {:ok, casted} ->
                      {:cont, {:ok, Map.put(record, load, casted)}}

                    :error ->
                      {:halt, {:error, message: "is invalid", field: calc.name}}
                  end
              end
          end
        end)

      other ->
        other
    end
  end

  def apply_opts_load(%Ash.Changeset{} = changeset, opts) do
    if opts[:load] do
      Ash.Changeset.load(changeset, opts[:load])
    else
      changeset
    end
  end

  def apply_opts_load(%Ash.Query{} = query, opts) do
    if opts[:load] do
      Ash.Query.load(query, opts[:load], Keyword.take(opts, [:strict?]))
    else
      query
    end
  end

  def load({:ok, result, instructions}, changeset, domain, opts) do
    if changeset.load in [nil, []] do
      {:ok, result, instructions}
    else
      query =
        changeset.resource
        |> Ash.Query.load(changeset.load)
        |> Ash.Query.set_context(changeset.context)
        |> select_selected(result)

      case Ash.load(result, query, Keyword.put(opts, :domain, domain)) do
        {:ok, result} ->
          {:ok, result, instructions}

        {:error, error} ->
          {:error, error}
      end
    end
  end

  def load({:ok, result}, changeset, domain, opts) do
    if changeset.load in [nil, []] do
      {:ok, result, %{}}
    else
      query =
        changeset.resource
        |> Ash.Query.load(changeset.load)
        |> Ash.Query.set_context(changeset.context)
        |> select_selected(result)

      case Ash.load(result, query, Keyword.put(opts, :domain, domain)) do
        {:ok, result} ->
          {:ok, result, %{}}

        {:error, error} ->
          {:error, error}
      end
    end
  end

  def load(other, _, _, _), do: other

  defp select_selected(query, result) do
    select =
      query.resource
      |> Ash.Resource.Info.attributes()
      |> Enum.filter(&Ash.Resource.selected?(result, &1.name))
      |> Enum.map(& &1.name)

    Ash.Query.ensure_selected(query, select)
  end

  def restrict_field_access(result, %Ash.Query{
        context: %{private: %{loading_relationship?: true}}
      }) do
    result
  end

  def restrict_field_access({:ok, record, instructions}, query_or_changeset) do
    {:ok, restrict_field_access(record, query_or_changeset), instructions}
  end

  def restrict_field_access({:ok, record}, query_or_changeset) do
    {:ok, restrict_field_access(record, query_or_changeset)}
  end

  def restrict_field_access({:error, error}, _), do: {:error, error}

  def restrict_field_access(records, query_or_changeset) when is_list(records) do
    Enum.map(records, &restrict_field_access(&1, query_or_changeset))
  end

  def restrict_field_access(%struct{results: results} = page, query_or_changeset)
      when struct in [Ash.Page.Keyset, Ash.Page.Offset] do
    %{page | results: restrict_field_access(results, query_or_changeset)}
  end

  def restrict_field_access(%Ash.NotLoaded{} = not_loaded, _query_or_changeset) do
    not_loaded
  end

  def restrict_field_access(%Ash.ForbiddenField{} = forbidden_field, _query_or_changeset) do
    forbidden_field
  end

  def restrict_field_access(%_{} = record, query_or_changeset) do
    embedded? = Ash.Resource.Info.embedded?(query_or_changeset.resource)

    if internal?(query_or_changeset) ||
         (embedded? && !query_or_changeset.context[:private][:cleaning_up_field_auth?]) do
      record
    else
      record.calculations
      |> Enum.reduce(record, fn
        {{:__ash_fields_are_visible__, fields}, value}, record ->
          if value do
            record
          else
            Enum.reduce(fields, record, fn field, record ->
              type =
                case Ash.Resource.Info.field(query_or_changeset.resource, field) do
                  %Ash.Resource.Aggregate{} -> :aggregate
                  %Ash.Resource.Attribute{} -> :attribute
                  %Ash.Resource.Calculation{} -> :calculation
                end

              forbidden_field =
                if embedded? && type == :attribute do
                  %Ash.ForbiddenField{
                    field: field,
                    type: type,
                    original_value: Map.get(record, field)
                  }
                else
                  %Ash.ForbiddenField{field: field, type: type}
                end

              record
              |> Map.put(field, forbidden_field)
              |> replace_dynamic_loads(field, type, query_or_changeset)
            end)
          end
          |> Map.update!(
            :calculations,
            &Map.delete(&1, {:__ash_fields_are_visible__, fields})
          )

        _, record ->
          record
      end)
    end
  end

  defp replace_dynamic_loads(record, _, :aggregate, _), do: record

  defp replace_dynamic_loads(record, field, type, %Ash.Changeset{} = changeset)
       when type in [:attribute, :calculation] do
    query =
      changeset.resource
      |> Ash.Query.new()
      |> Ash.Query.load(changeset.load)

    replace_dynamic_loads(record, field, type, query)
  end

  defp replace_dynamic_loads(record, field, type, query)
       when type in [:attribute, :calculation] do
    Enum.reduce(
      query.calculations,
      record,
      fn
        {key, %{module: Ash.Resource.Calculation.LoadAttribute, opts: opts, load: load}},
        record ->
          if type == :attribute && opts[:attribute] == field do
            if load do
              Map.put(record, load, %Ash.ForbiddenField{field: load, type: type})
            else
              Map.update!(
                record,
                :calculations,
                &Map.put(&1, key, %Ash.ForbiddenField{field: field, type: type})
              )
            end
          else
            record
          end

        {key, %{calc_name: calc_name, load: load}}, record ->
          if calc_name == field and type == :calculation do
            if load do
              Map.put(record, load, %Ash.ForbiddenField{field: load, type: type})
            else
              Map.update!(
                record,
                :calculations,
                &Map.put(&1, key, %Ash.ForbiddenField{field: field, type: type})
              )
            end
          else
            record
          end

        _, record ->
          record
      end
    )
  end

  def select({:ok, results, instructions}, query) do
    {:ok, select(results, query), instructions}
  end

  def select({:ok, results}, query) do
    {:ok, select(results, query)}
  end

  def select({:error, error}, _query) do
    {:error, error}
  end

  def select(nil, _), do: nil

  def select(result, %{select: nil}) do
    result
  end

  def select(result, nil) do
    result
  end

  def select(%resource{} = result, %{select: select, resource: resource} = query) do
    select_mask = select_mask(query)

    result
    |> Map.merge(select_mask)
    |> Ash.Resource.put_metadata(:selected, select)
  end

  def select(:ok, _query), do: :ok

  def select(results, %{select: select} = query) do
    if Enumerable.impl_for(results) do
      select_mask = select_mask(query)

      Enum.map(results, fn result ->
        result
        |> Map.merge(select_mask)
        |> Ash.Resource.put_metadata(:selected, select)
      end)
    else
      results
    end
  end

  defp select_mask(%{select: select, resource: resource}) do
    resource
    |> Ash.Resource.Info.attributes()
    |> Enum.reject(fn attribute ->
      if is_nil(select) do
        attribute.select_by_default?
      else
        attribute.always_select? || attribute.primary_key? || attribute.name in select
      end
    end)
    |> Map.new(fn attribute ->
      {attribute.name, %Ash.NotLoaded{field: attribute.name, type: :attribute}}
    end)
  end
end
