# SPDX-FileCopyrightText: 2022 Alembic Pty Ltd
#
# SPDX-License-Identifier: MIT

# credo:disable-for-this-file Credo.Check.Design.AliasUsage
if Code.ensure_loaded?(Igniter) do
  defmodule Mix.Tasks.AshAuthentication.Upgrade do
    @moduledoc false

    use Igniter.Mix.Task

    @impl Igniter.Mix.Task
    def info(_argv, _composing_task) do
      %Igniter.Mix.Task.Info{
        # Groups allow for overlapping arguments for tasks by the same author
        # See the generators guide for more.
        group: :ash_authentication,
        # *other* dependencies to add
        # i.e `{:foo, "~> 2.0"}`
        adds_deps: [],
        # *other* dependencies to add and call their associated installers, if they exist
        # i.e `{:foo, "~> 2.0"}`
        installs: [],
        # An example invocation
        # example: __MODULE__.Docs.example(),
        example: "example",
        # a list of positional arguments, i.e `[:file]`
        positional: [:from, :to],
        # Other tasks your task composes using `Igniter.compose_task`, passing in the CLI argv
        # This ensures your option schema includes options from nested tasks
        composes: [],
        # `OptionParser` schema
        schema: [],
        # Default values for the options in the `schema`
        defaults: [],
        # CLI aliases
        aliases: [],
        # A list of options in the schema that are required
        required: []
      }
    end

    @impl Igniter.Mix.Task
    def igniter(igniter) do
      positional = igniter.args.positional
      options = igniter.args.options

      upgrades =
        %{
          "4.4.9" => [&fix_token_is_revoked_action/2],
          "4.13.4" => [&add_remember_me_to_magic_link_sign_in/2],
          "5.0.0" => [
            &fix_google_hd_field/2,
            &convert_revoked_read_action_to_generic/2,
            &convert_request_actions_to_generic/2,
            &add_brute_force_protection/2,
            &require_identity_resource/2
          ]
        }

      # For each version that requires a change, add it to this map
      # Each key is a version that points at a list of functions that take an
      # igniter and options (i.e. flags or other custom options).
      # See the upgrades guide for more.
      Igniter.Upgrades.run(igniter, positional.from, positional.to, upgrades,
        custom_opts: options
      )
    end

    def fix_token_is_revoked_action(igniter, _opts) do
      case find_all_token_resources(igniter) do
        {igniter, []} ->
          igniter

        {igniter, resources} ->
          Enum.reduce(resources, igniter, fn resource, igniter ->
            maybe_fix_is_revoked_action(igniter, resource)
          end)
      end
    end

    @doc """
    Fixes the Google strategy field name change from "google_hd" to "hd".

    The Google strategy now uses OIDC which returns standard
    claims. The hosted domain claim changed from "google_hd" to "hd".
    """
    def fix_google_hd_field(igniter, _opts) do
      igniter
      |> replace_google_hd_strings()
      |> Igniter.add_notice("""
      Google Strategy Breaking Change:

      The `email_verified` field in `user_info` is now a `boolean` instead of a `string`.

      If your code checks `user_info["email_verified"] == "true"`, update it to:
        user_info["email_verified"] == true

      Please review your `register_with_google` action and any code that accesses
      the `email_verified` field from Google OAuth responses.
      """)
    end

    defp replace_google_hd_strings(igniter) do
      igniter = Igniter.include_all_elixir_files(igniter)

      igniter.rewrite
      |> Rewrite.sources()
      |> Enum.filter(&match?(%Rewrite.Source{filetype: %Rewrite.Source.Ex{}}, &1))
      |> Enum.reduce(igniter, fn source, igniter ->
        zipper =
          source
          |> Rewrite.Source.get(:quoted)
          |> Sourceror.Zipper.zip()

        case replace_google_hd_in_zipper(zipper) do
          {:ok, new_zipper} ->
            new_quoted = Sourceror.Zipper.topmost_root(new_zipper)
            new_source = Igniter.update_source(source, igniter, :quoted, new_quoted)
            %{igniter | rewrite: Rewrite.update!(igniter.rewrite, new_source)}

          :unchanged ->
            igniter
        end
      end)
    end

    defp replace_google_hd_in_zipper(zipper) do
      {new_zipper, changed?} =
        Sourceror.Zipper.traverse(zipper, false, fn zipper, changed? ->
          if google_hd_string?(zipper) do
            {replace_with_hd(zipper), true}
          else
            {zipper, changed?}
          end
        end)

      if changed? do
        {:ok, new_zipper}
      else
        :unchanged
      end
    end

    defp google_hd_string?(%{node: "google_hd"}), do: true

    defp google_hd_string?(%{node: {:__block__, meta, ["google_hd"]}}) when is_list(meta),
      do: true

    defp google_hd_string?(_), do: false

    defp replace_with_hd(%{node: "google_hd"} = zipper) do
      Sourceror.Zipper.replace(zipper, "hd")
    end

    defp replace_with_hd(%{node: {:__block__, meta, ["google_hd"]}} = zipper) do
      Sourceror.Zipper.replace(zipper, {:__block__, meta, ["hd"]})
    end

    defp find_all_token_resources(igniter) do
      Igniter.Project.Module.find_all_matching_modules(igniter, fn _module, zipper ->
        with {:ok, zipper} <- Igniter.Code.Module.move_to_use(zipper, Ash.Resource),
             {:ok, zipper} <- Igniter.Code.Function.move_to_nth_argument(zipper, 1),
             {:ok, zipper} <- Igniter.Code.Keyword.get_key(zipper, :extensions) do
          extensions_contain_token_resource?(zipper)
        end
      end)
    end

    defp extensions_contain_token_resource?(zipper) do
      if Igniter.Code.List.list?(zipper) do
        match?(
          {:ok, _},
          Igniter.Code.List.move_to_list_item(
            zipper,
            &Igniter.Code.Common.nodes_equal?(&1, AshAuthentication.TokenResource)
          )
        )
      else
        Igniter.Code.Common.nodes_equal?(zipper, AshAuthentication.TokenResource)
      end
    end

    defp maybe_fix_is_revoked_action(igniter, resource) do
      Igniter.Project.Module.find_and_update_module!(igniter, resource, fn zipper ->
        with {:ok, action_zipper} <- move_to_action(zipper, :action, :revoked?),
             {:ok, zipper} <- Igniter.Code.Common.move_to_do_block(action_zipper),
             {:ok, zipper} <- remove_argument_option(zipper, :token, :allow_nil?),
             {:ok, zipper} <- remove_argument_option(zipper, :jti, :allow_nil?) do
          add_action_return_type(zipper, :boolean)
        else
          :error -> {:ok, zipper}
        end
      end)
    end

    defp move_to_action(zipper, type, name) do
      Igniter.Code.Function.move_to_function_call(
        zipper,
        type,
        2,
        &Igniter.Code.Function.argument_equals?(&1, 0, name)
      )
    end

    defp add_action_return_type(zipper, type) do
      zipper = Sourceror.Zipper.top(zipper)

      with {:ok, zipper} <- move_to_action(zipper, :action, :revoked?),
           {:ok, zipper} <- Igniter.Code.Function.move_to_nth_argument(zipper, 1) do
        {:ok,
         Sourceror.Zipper.insert_left(
           zipper,
           quote do
             unquote(type)
           end
         )}
      end
    end

    defp remove_argument_option(zipper, argument_name, option) do
      with {:ok, zipper} <-
             Igniter.Code.Function.move_to_function_call(
               zipper,
               :argument,
               3,
               &Igniter.Code.Function.argument_equals?(&1, 0, argument_name)
             ),
           {:ok, zipper} <- Igniter.Code.Function.move_to_nth_argument(zipper, 2) do
        remove_option_from_argument(zipper, option)
      end
    end

    defp remove_option_from_argument(zipper, option) do
      if Igniter.Code.List.find_list_item_index(
           zipper,
           &Igniter.Code.Tuple.elem_equals?(&1, 0, :do)
         ) do
        Igniter.Code.Common.within(zipper, fn zipper ->
          {:ok,
           Igniter.Code.Common.remove(
             zipper,
             &Igniter.Code.Function.function_call?(&1, option, 1)
           )}
        end)
      else
        Igniter.Code.List.remove_from_list(
          zipper,
          &Igniter.Code.Tuple.elem_equals?(&1, 0, option)
        )
      end
    end

    def add_remember_me_to_magic_link_sign_in(igniter, _opts) do
      case find_resources_with_magic_link_and_remember_me(igniter) do
        {igniter, []} ->
          igniter

        {igniter, resources} ->
          Enum.reduce(resources, igniter, fn resource, igniter ->
            maybe_add_remember_me_to_magic_link_action(igniter, resource)
          end)
      end
    end

    defp find_resources_with_magic_link_and_remember_me(igniter) do
      Igniter.Project.Module.find_all_matching_modules(igniter, fn _module, zipper ->
        with {:ok, zipper} <- enter_auth_strategies(zipper),
             true <- has_strategy?(zipper, :magic_link),
             true <- has_strategy?(zipper, :remember_me) do
          true
        else
          _ -> false
        end
      end)
    end

    @oauth2_family ~w[oauth2 oidc github google auth0 apple slack microsoft okta]a

    def require_identity_resource(igniter, _opts) do
      case find_resources_with_oauth2_strategies(igniter) do
        {igniter, []} ->
          igniter

        {igniter, resources} ->
          resources
          |> Enum.reduce(igniter, &ensure_identity_resource/2)
          |> Igniter.add_notice("""
          The user identity resource's unique key changed from
          `(strategy, uid, user_id)` to `(strategy, uid)`, so that a provider's
          `iss`/`sub` resolves to exactly one local user.

          Run `mix ash.codegen require_user_identity_unique_key` (and then
          `mix ash.migrate`) to generate the migration that swaps the unique
          index.

          IMPORTANT: the new index will fail to create if your data contains the
          same `(strategy, uid)` linked to more than one user. That should not
          happen under normal use, but if it does you must reconcile those rows
          before migrating - it indicates a provider identity was linked to
          multiple accounts.
          """)
      end
    end

    defp find_resources_with_oauth2_strategies(igniter) do
      Igniter.Project.Module.find_all_matching_modules(igniter, fn _module, zipper ->
        case enter_auth_strategies(zipper) do
          {:ok, zipper} -> Enum.any?(@oauth2_family, &has_strategy?(zipper, &1))
          _ -> false
        end
      end)
    end

    defp ensure_identity_resource(resource, igniter) do
      identity_resource = conventional_identity_resource(resource)

      case Igniter.Project.Module.module_exists(igniter, identity_resource) do
        {true, igniter} ->
          Enum.reduce(
            @oauth2_family,
            igniter,
            &wire_identity_resource(&2, resource, &1, identity_resource)
          )

        {false, igniter} ->
          Igniter.add_warning(
            igniter,
            missing_identity_resource_warning(resource, identity_resource)
          )
      end
    end

    defp wire_identity_resource(igniter, resource, type, identity_resource) do
      case AshAuthentication.Igniter.defines_strategy_of_type(igniter, resource, type) do
        {igniter, true} ->
          igniter
          |> add_identity_resource_to_strategy(resource, type, identity_resource)
          |> ensure_register_action_has_identity_change(resource, type)

        {igniter, false} ->
          igniter
      end
    end

    defp conventional_identity_resource(resource) do
      resource
      |> Module.split()
      |> :lists.droplast()
      |> Enum.concat(["UserIdentity"])
      |> Module.concat()
    end

    defp add_identity_resource_to_strategy(igniter, resource, type, identity_resource) do
      Igniter.Project.Module.find_and_update_module!(igniter, resource, fn zipper ->
        with {:ok, zipper} <- enter_auth_strategies(zipper),
             {:ok, strategy_zipper} <-
               Igniter.Code.Function.move_to_function_call_in_current_scope(zipper, type, [1, 2]),
             {:ok, do_block_zipper} <- Igniter.Code.Common.move_to_do_block(strategy_zipper),
             :error <-
               Igniter.Code.Function.move_to_function_call_in_current_scope(
                 do_block_zipper,
                 :identity_resource,
                 1
               ) do
          {:ok,
           Igniter.Code.Common.add_code(
             do_block_zipper,
             "identity_resource #{inspect(identity_resource)}"
           )}
        else
          _ -> {:ok, zipper}
        end
      end)
    end

    # sobelow_skip ["DOS.BinToAtom"]
    defp ensure_register_action_has_identity_change(igniter, resource, type) do
      action_name = :"register_with_#{type}"

      Igniter.Project.Module.find_and_update_module!(igniter, resource, fn zipper ->
        with {:ok, action_zipper} <- move_to_action(zipper, :create, action_name),
             {:ok, do_block_zipper} <- Igniter.Code.Common.move_to_do_block(action_zipper),
             false <- has_identity_change?(do_block_zipper) do
          {:ok,
           Igniter.Code.Common.add_code(
             do_block_zipper,
             "change AshAuthentication.Strategy.OAuth2.IdentityChange"
           )}
        else
          _ -> {:ok, zipper}
        end
      end)
    end

    defp has_identity_change?(zipper) do
      match?(
        {:ok, _},
        Igniter.Code.Function.move_to_function_call_in_current_scope(
          zipper,
          :change,
          1,
          fn change_zipper ->
            case Igniter.Code.Function.move_to_nth_argument(change_zipper, 0) do
              {:ok, arg_zipper} ->
                arg_zipper
                |> Sourceror.Zipper.node()
                |> Sourceror.to_string()
                |> String.contains?("IdentityChange")

              _ ->
                false
            end
          end
        )
      )
    end

    defp missing_identity_resource_warning(resource, identity_resource) do
      """
      #{inspect(resource)} has one or more OAuth2/OIDC strategies but no user
      identity resource could be found at #{inspect(identity_resource)}.

      As of this release, OAuth2 and OIDC strategies require an `identity_resource`.
      Matching a local user by their email address (or any other provider claim)
      is unsafe - only the provider's `iss`/`sub` claims uniquely and stably
      identify an end-user, and those are persisted in the identity resource.

      To resolve this manually:

        1. Create a user identity resource (conventionally #{inspect(identity_resource)}):

           defmodule #{inspect(identity_resource)} do
             use Ash.Resource,
               extensions: [AshAuthentication.UserIdentity],
               domain: <your domain>

             user_identity do
               user_resource #{inspect(resource)}
             end

             # ... data layer, postgres/sqlite block, etc.
           end

        2. Add `identity_resource #{inspect(identity_resource)}` to each OAuth2/OIDC
           strategy on #{inspect(resource)}.

        3. Add `change AshAuthentication.Strategy.OAuth2.IdentityChange` to each
           `register_with_*` action for those strategies.

      See the "User Identities" section of the strategy documentation for details.
      """
    end

    defp enter_auth_strategies(zipper) do
      with {:ok, zipper} <-
             Igniter.Code.Function.move_to_function_call_in_current_scope(
               zipper,
               :authentication,
               1
             ),
           {:ok, zipper} <- Igniter.Code.Common.move_to_do_block(zipper),
           {:ok, zipper} <-
             Igniter.Code.Function.move_to_function_call_in_current_scope(
               zipper,
               :strategies,
               1
             ) do
        Igniter.Code.Common.move_to_do_block(zipper)
      end
    end

    defp has_strategy?(zipper, strategy_type) do
      match?(
        {:ok, _},
        Igniter.Code.Function.move_to_function_call_in_current_scope(
          zipper,
          strategy_type,
          [1, 2]
        )
      )
    end

    defp maybe_add_remember_me_to_magic_link_action(igniter, resource) do
      Igniter.Project.Module.find_and_update_module!(igniter, resource, fn zipper ->
        with {:ok, action_zipper} <- move_to_action(zipper, :create, :sign_in_with_magic_link),
             {:ok, do_block_zipper} <- Igniter.Code.Common.move_to_do_block(action_zipper) do
          do_block_zipper
          |> maybe_add_remember_me_argument()
          |> maybe_add_remember_me_change()
        else
          :error -> {:ok, zipper}
        end
      end)
    end

    defp maybe_add_remember_me_argument(zipper) do
      if has_remember_me_argument?(zipper) do
        zipper
      else
        add_remember_me_argument(zipper)
      end
    end

    defp has_remember_me_argument?(zipper) do
      match?(
        {:ok, _},
        Igniter.Code.Function.move_to_function_call_in_current_scope(
          zipper,
          :argument,
          [2, 3],
          &Igniter.Code.Function.argument_equals?(&1, 0, :remember_me)
        )
      )
    end

    defp add_remember_me_argument(zipper) do
      argument_code = """
      argument :remember_me, :boolean do
        description "Whether to generate a remember me token"
        allow_nil? true
      end
      """

      Igniter.Code.Common.add_code(zipper, argument_code)
    end

    defp maybe_add_remember_me_change(zipper) do
      if has_remember_me_change?(zipper) do
        {:ok, zipper}
      else
        {:ok, add_remember_me_change(zipper)}
      end
    end

    defp has_remember_me_change?(zipper) do
      match?(
        {:ok, _},
        Igniter.Code.Function.move_to_function_call_in_current_scope(
          zipper,
          :change,
          1,
          fn change_zipper ->
            case Igniter.Code.Function.move_to_nth_argument(change_zipper, 0) do
              {:ok, arg_zipper} ->
                source = Sourceror.Zipper.node(arg_zipper) |> Sourceror.to_string()
                String.contains?(source, "MaybeGenerateTokenChange")

              _ ->
                false
            end
          end
        )
      )
    end

    defp add_remember_me_change(zipper) do
      change_code = """
      change {AshAuthentication.Strategy.RememberMe.MaybeGenerateTokenChange,
              strategy_name: :remember_me}
      """

      Igniter.Code.Common.add_code(zipper, change_code)
    end

    def convert_revoked_read_action_to_generic(igniter, _opts) do
      case find_all_token_resources(igniter) do
        {igniter, []} ->
          igniter

        {igniter, resources} ->
          Enum.reduce(resources, igniter, fn resource, igniter ->
            maybe_convert_revoked_read_to_generic(igniter, resource)
          end)
      end
    end

    @doc """
    Converts request actions from :read to :action type.

    In AshAuthentication 5.0, both password reset request and magic link request
    actions were changed from :read to :action type. This upgrade function helps
    migrate existing custom actions.
    """
    def convert_request_actions_to_generic(igniter, _opts) do
      igniter
      |> convert_password_reset_request_actions()
      |> convert_magic_link_request_actions()
    end

    defp convert_password_reset_request_actions(igniter) do
      case find_resources_with_password_resettable(igniter) do
        {igniter, []} ->
          igniter

        {igniter, resources} ->
          Enum.reduce(resources, igniter, fn resource, igniter ->
            convert_password_reset_request_action(igniter, resource)
          end)
      end
    end

    defp maybe_convert_revoked_read_to_generic(igniter, resource) do
      Igniter.Project.Module.find_and_update_module!(igniter, resource, fn zipper ->
        with {:ok, action_zipper} <- move_to_action(zipper, :read, :revoked?),
             {:ok, do_block_zipper} <- Igniter.Code.Common.move_to_do_block(action_zipper) do
          convert_read_to_generic_action(do_block_zipper)
        else
          :error -> {:ok, zipper}
        end
      end)
    end

    defp find_resources_with_password_resettable(igniter) do
      Igniter.Project.Module.find_all_matching_modules(igniter, fn _module, zipper ->
        with {:ok, zipper} <- enter_auth_strategies(zipper),
             true <- has_strategy?(zipper, :password) do
          password_has_resettable?(zipper)
        else
          _ -> false
        end
      end)
    end

    defp password_has_resettable?(zipper) do
      with {:ok, password_zipper} <-
             Igniter.Code.Function.move_to_function_call_in_current_scope(
               zipper,
               :password,
               [1, 2]
             ),
           {:ok, do_block} <- Igniter.Code.Common.move_to_do_block(password_zipper) do
        match?(
          {:ok, _},
          Igniter.Code.Function.move_to_function_call_in_current_scope(
            do_block,
            :resettable,
            1
          )
        )
      else
        _ -> false
      end
    end

    defp convert_password_reset_request_action(igniter, resource) do
      Igniter.Project.Module.find_and_update_module!(igniter, resource, fn zipper ->
        with {:ok, action_zipper} <-
               move_to_action(zipper, :read, :request_password_reset_with_password),
             {:ok, do_block} <- Igniter.Code.Common.move_to_do_block(action_zipper) do
          identity_field = get_identity_field_from_action(do_block) || :email

          zipper
          |> ensure_get_by_action_exists(identity_field)
          |> convert_read_action_to_generic(
            :request_password_reset_with_password,
            identity_field,
            :password
          )
        else
          :error -> {:ok, zipper}
        end
      end)
    end

    defp convert_read_to_generic_action(do_block_zipper) do
      do_block_zipper =
        Igniter.Code.Common.remove(
          do_block_zipper,
          &Igniter.Code.Function.function_call?(&1, :get?, 1)
        )

      do_block_zipper =
        Igniter.Code.Common.remove(
          do_block_zipper,
          &Igniter.Code.Function.function_call?(&1, :prepare, 1)
        )

      do_block_zipper =
        Igniter.Code.Common.add_code(
          do_block_zipper,
          "run AshAuthentication.TokenResource.IsRevoked"
        )

      zipper = Sourceror.Zipper.top(do_block_zipper)

      with {:ok, zipper} <- move_to_action(zipper, :read, :revoked?) do
        new_node =
          Sourceror.Zipper.node(zipper)
          |> replace_read_with_action()
          |> insert_boolean_return_type()

        {:ok, Sourceror.Zipper.replace(zipper, new_node)}
      end
    end

    defp replace_read_with_action({:read, meta, args}) do
      {:action, meta, args}
    end

    defp insert_boolean_return_type({:action, meta, [name | rest]}) do
      {:action, meta, [name, :boolean | rest]}
    end

    defp convert_magic_link_request_actions(igniter) do
      case find_resources_with_magic_link(igniter) do
        {igniter, []} ->
          igniter

        {igniter, resources} ->
          Enum.reduce(resources, igniter, fn resource, igniter ->
            convert_magic_link_request_action(igniter, resource)
          end)
      end
    end

    defp find_resources_with_magic_link(igniter) do
      Igniter.Project.Module.find_all_matching_modules(igniter, fn _module, zipper ->
        case enter_auth_strategies(zipper) do
          {:ok, zipper} -> has_strategy?(zipper, :magic_link)
          _ -> false
        end
      end)
    end

    defp convert_magic_link_request_action(igniter, resource) do
      Igniter.Project.Module.find_and_update_module!(igniter, resource, fn zipper ->
        with {:ok, action_zipper} <- move_to_action(zipper, :read, :request_magic_link),
             {:ok, do_block} <- Igniter.Code.Common.move_to_do_block(action_zipper) do
          identity_field = get_identity_field_from_action(do_block) || :email

          zipper
          |> ensure_get_by_action_exists(identity_field)
          |> convert_read_action_to_generic(:request_magic_link, identity_field, :magic_link)
        else
          :error -> {:ok, zipper}
        end
      end)
    end

    defp get_identity_field_from_action(do_block) do
      with {:ok, arg_zipper} <-
             Igniter.Code.Function.move_to_function_call_in_current_scope(
               do_block,
               :argument,
               [2, 3]
             ),
           {:ok, name_zipper} <- Igniter.Code.Function.move_to_nth_argument(arg_zipper, 0) do
        case Sourceror.Zipper.node(name_zipper) do
          name when is_atom(name) -> name
          _ -> nil
        end
      else
        _ -> nil
      end
    end

    defp ensure_get_by_action_exists(zipper, identity_field) do
      action_name = :"get_by_#{identity_field}"

      if action_exists?(zipper, :read, action_name) do
        zipper
      else
        action_code = """
        read :#{action_name} do
          get_by :#{identity_field}
          description "Look up a user by #{identity_field}."
        end
        """

        with {:ok, actions_zipper} <- move_to_actions_block(zipper),
             {:ok, do_block} <- Igniter.Code.Common.move_to_do_block(actions_zipper) do
          Igniter.Code.Common.add_code(do_block, action_code)
          |> Sourceror.Zipper.top()
        else
          _ -> zipper
        end
      end
    end

    defp action_exists?(zipper, type, name) do
      match?({:ok, _}, move_to_action(zipper, type, name))
    end

    defp move_to_actions_block(zipper) do
      Igniter.Code.Function.move_to_function_call_in_current_scope(zipper, :actions, 1)
    end

    defp convert_read_action_to_generic(zipper, action_name, identity_field, strategy_type) do
      run_module =
        case strategy_type do
          :password -> "AshAuthentication.Strategy.Password.RequestPasswordReset"
          :magic_link -> "AshAuthentication.Strategy.MagicLink.Request"
        end

      new_action_code = """
      action :#{action_name} do
        argument :#{identity_field}, :ci_string, allow_nil?: false
        run {#{run_module}, action: :get_by_#{identity_field}}
        description "Send #{strategy_type_description(strategy_type)} to a user if they exist."
      end
      """

      case move_to_action(zipper, :read, action_name) do
        {:ok, action_zipper} ->
          action_zipper
          |> Sourceror.Zipper.remove()
          |> Sourceror.Zipper.top()
          |> replace_with_generic_action(new_action_code)
          |> then(&{:ok, &1})

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

    defp replace_with_generic_action(zipper, new_action_code) do
      with {:ok, actions_zipper} <- move_to_actions_block(zipper),
           {:ok, do_block} <- Igniter.Code.Common.move_to_do_block(actions_zipper) do
        Igniter.Code.Common.add_code(do_block, new_action_code)
        |> Sourceror.Zipper.top()
      else
        _ -> zipper
      end
    end

    defp strategy_type_description(:password), do: "password reset instructions"
    defp strategy_type_description(:magic_link), do: "a magic link"

    @doc """
    Adds identity-keyed brute-force protection to password and magic link strategies.

    For each resource that uses the `password` or `magic_link` strategy, this
    upgrade ensures an `audit_log` add-on is present (composing the
    `ash_authentication.add_add_on.audit_log` task to generate one if needed)
    and adds `brute_force_strategy {:audit_log, <name>}` to each affected
    strategy block.
    """
    def add_brute_force_protection(igniter, _opts) do
      case find_resources_needing_brute_force(igniter) do
        {igniter, []} ->
          igniter

        {igniter, resources_with_strategies} ->
          resources_with_strategies
          |> Enum.reduce(igniter, &add_brute_force_to_resource/2)
          |> Igniter.add_notice(brute_force_notice())
      end
    end

    defp add_brute_force_to_resource({resource, strategies}, igniter) do
      {igniter, audit_log_name} = ensure_audit_log(igniter, resource)

      Enum.reduce(strategies, igniter, fn strategy_type, igniter ->
        add_brute_force_to_strategy(igniter, resource, strategy_type, audit_log_name)
      end)
    end

    defp brute_force_notice do
      """
      Brute-force Protection:

      Identity-keyed brute-force protection has been added to your password
      and magic link strategies via the audit log add-on. Failed sign-in,
      password reset request, and magic link request attempts are now
      counted within a 5 minute window per identity, with requests blocked
      after 5 failures.

      Adjust `audit_log_window` and `audit_log_max_failures` on each
      strategy if you need different thresholds. If your audit log uses
      `exclude_actions` or `exclude_strategies` to skip the protected
      actions, the compile-time verifier will report which actions need
      to be tracked.
      """
    end

    defp find_resources_needing_brute_force(igniter) do
      {igniter, resources} = find_resources_with_password_or_magic_link(igniter)
      Enum.reduce(resources, {igniter, []}, &collect_strategies_needing_brute_force/2)
    end

    defp collect_strategies_needing_brute_force(resource, {igniter, acc}) do
      {igniter, strategies_needing} =
        Enum.reduce(
          [:password, :magic_link],
          {igniter, []},
          &reduce_strategy_needing_brute_force(&1, &2, resource)
        )

      case strategies_needing do
        [] -> {igniter, acc}
        strategies -> {igniter, [{resource, Enum.reverse(strategies)} | acc]}
      end
    end

    defp reduce_strategy_needing_brute_force(strategy_type, {igniter, strategies}, resource) do
      case strategy_needs_brute_force?(igniter, resource, strategy_type) do
        {igniter, true} -> {igniter, [strategy_type | strategies]}
        {igniter, false} -> {igniter, strategies}
      end
    end

    defp strategy_needs_brute_force?(igniter, resource, strategy_type) do
      result =
        Spark.Igniter.find(igniter, resource, fn _, zipper ->
          check_strategy_for_brute_force(zipper, strategy_type)
        end)

      case result do
        {:ok, igniter, _module, true} -> {igniter, true}
        {:error, igniter} -> {igniter, false}
      end
    end

    defp check_strategy_for_brute_force(zipper, strategy_type) do
      with {:ok, strategies_zipper} <- enter_auth_strategies(zipper),
           {:ok, strategy_zipper} <-
             Igniter.Code.Function.move_to_function_call_in_current_scope(
               strategies_zipper,
               strategy_type,
               [1, 2]
             ),
           {:ok, do_block_zipper} <- Igniter.Code.Common.move_to_do_block(strategy_zipper),
           false <- has_brute_force_strategy?(do_block_zipper) do
        {:ok, true}
      else
        _ -> :error
      end
    end

    defp find_resources_with_password_or_magic_link(igniter) do
      Igniter.Project.Module.find_all_matching_modules(igniter, fn _module, zipper ->
        case enter_auth_strategies(zipper) do
          {:ok, zipper} ->
            has_strategy?(zipper, :password) or has_strategy?(zipper, :magic_link)

          _ ->
            false
        end
      end)
    end

    defp ensure_audit_log(igniter, resource) do
      case find_audit_log_name(igniter, resource) do
        {igniter, nil} ->
          igniter =
            Igniter.compose_task(
              igniter,
              "ash_authentication.add_add_on.audit_log",
              ["--user", inspect(resource)]
            )

          {igniter, :audit_log}

        {igniter, name} ->
          {igniter, name}
      end
    end

    defp find_audit_log_name(igniter, resource) do
      result =
        Spark.Igniter.find(igniter, resource, fn _, zipper ->
          with {:ok, zipper} <-
                 Igniter.Code.Function.move_to_function_call_in_current_scope(
                   zipper,
                   :authentication,
                   1
                 ),
               {:ok, zipper} <- Igniter.Code.Common.move_to_do_block(zipper),
               {:ok, zipper} <-
                 Igniter.Code.Function.move_to_function_call_in_current_scope(
                   zipper,
                   :add_ons,
                   1
                 ),
               {:ok, zipper} <- Igniter.Code.Common.move_to_do_block(zipper),
               {:ok, audit_log_zipper} <-
                 Igniter.Code.Function.move_to_function_call_in_current_scope(
                   zipper,
                   :audit_log,
                   [1, 2]
                 ) do
            {:ok, audit_log_name(audit_log_zipper)}
          else
            _ -> :error
          end
        end)

      case result do
        {:ok, igniter, _module, name} -> {igniter, name}
        {:error, igniter} -> {igniter, nil}
      end
    end

    defp audit_log_name(zipper) do
      case Igniter.Code.Function.move_to_nth_argument(zipper, 0) do
        {:ok, name_zipper} ->
          case Sourceror.Zipper.node(name_zipper) do
            name when is_atom(name) -> name
            {:__block__, _, [name]} when is_atom(name) -> name
            _ -> :audit_log
          end

        _ ->
          :audit_log
      end
    end

    defp add_brute_force_to_strategy(igniter, resource, strategy_type, audit_log_name) do
      Igniter.Project.Module.find_and_update_module!(igniter, resource, fn zipper ->
        with {:ok, strategies_zipper} <- enter_auth_strategies(zipper),
             {:ok, strategy_zipper} <-
               Igniter.Code.Function.move_to_function_call_in_current_scope(
                 strategies_zipper,
                 strategy_type,
                 [1, 2]
               ),
             {:ok, do_block_zipper} <- Igniter.Code.Common.move_to_do_block(strategy_zipper) do
          updated_zipper =
            Igniter.Code.Common.add_code(
              do_block_zipper,
              "brute_force_strategy {:audit_log, #{inspect(audit_log_name)}}"
            )

          {:ok, Sourceror.Zipper.top(updated_zipper)}
        else
          _ -> {:ok, zipper}
        end
      end)
    end

    defp has_brute_force_strategy?(do_block_zipper) do
      match?(
        {:ok, _},
        Igniter.Code.Function.move_to_function_call_in_current_scope(
          do_block_zipper,
          :brute_force_strategy,
          1
        )
      )
    end
  end
else
  defmodule Mix.Tasks.AshAuthentication.Upgrade do
    @moduledoc false

    use Mix.Task

    def run(_argv) do
      Mix.shell().error("""
      The task 'ash_authentication.upgrade' requires igniter. Please install igniter and try again.

      For more information, see: https://hexdocs.pm/igniter/readme.html#installation
      """)

      exit({:shutdown, 1})
    end
  end
end