defmodule Toolbox.Utils.Enum do
  @moduledoc """
  A set of utility functions for `Enumerable`s.
  """

  @doc """
  Returns `:ok` if `fun` applied to every element of `enumerable` returns `:ok`.
  Otherwise it returns the first error occurred during `fun` call.

  ## Example
      iex> each_while_ok([], fn _ -> :not_called_at_all end)
      :ok

      iex> each_while_ok([1, 2, 3], fn _ -> :ok end)
      :ok

      iex> each_while_ok([1, 2, 3], &(if &1 == 2, do: {:error, &1}, else: :ok))
      {:error, 2}
  """
  @spec each_while_ok(Enumerable.t(), (any() -> result)) :: result
        when result: :ok | {:error, reason :: any()}
  def each_while_ok(enumerable, fun) do
    with {:ok, :ok} <- reduce_while_ok(enumerable, :ok, &each_if_ok_reducer(&1, &2, fun)) do
      :ok
    end
  end

  @doc """
  Returns `{:ok, list_of_values}` if `fun` applied to every element of
  `enumerable` returns `{:ok, value}`. Otherwise it returns the first
  error occurred during `fun` call.

  ## Example
      iex> map_while_ok([], fn _ -> :not_called_at_all end)
      {:ok, []}

      iex> map_while_ok([1, 2, 3], &({:ok, &1 * 2}))
      {:ok, [2, 4, 6]}

      iex> map_while_ok([1, 2, 3], &(if &1 == 2, do: {:error, &1 * 2}, else: {:ok, &1 * 2}))
      {:error, 4}
  """
  @spec map_while_ok(Enumerable.t(), (any() -> {:ok, any()} | {:error, reason :: any()})) ::
          {:ok, list()} | {:error, reason :: any()}
  def map_while_ok(enumerable, fun) do
    with {:ok, items} <- reduce_while_ok(enumerable, [], &map_if_ok_reducer(&1, &2, fun)) do
      {:ok, Enum.reverse(items)}
    end
  end

  @doc """
  Invokes `fun` for each element in the `enumerable` with the accumulator.
  If every `fun` call returns `{:ok, new_acc}`, functions returns `{:ok, last_acc}`.
  Otherwise it returns the first error occurred during `fun` call.

  ## Example
      iex> reduce_while_ok([], :initial, fn _, _ -> :not_called_at_all end)
      {:ok, :initial}

      iex> reduce_while_ok([1, 2, 3], 0, &({:ok, &1 + &2}))
      {:ok, 6}

      iex> reduce_while_ok([1, 2, 3], 0, &(if &1 == 2, do: {:error, &1 + &2}, else: {:ok, &1 + &2}))
      {:error, 3}
  """
  @spec reduce_while_ok(
          Enumerable.t(),
          any(),
          (any(), any() -> {:ok, any()} | {:error, reason: any()})
        ) :: {:ok, any()} | {:error, reason :: any()}
  def reduce_while_ok(enumerable, acc, fun) do
    Enum.reduce_while(enumerable, {:ok, acc}, &halt_on_error(&1, &2, fun))
  end

  @doc """
  Flattens the enum by the specified amount of levels.

  Flattens only the specified amount of levels, contrary to the standard `List.flatten/2` which
  flattens the list recursively. Handy if you e.g. want to flatten only the first level.

  Importantly, this operation maintains order of the elements in the list.

  ## Examples
      iex> flatten_levels([[:hello, :team], ["this", "is"], ["now", ["flat", ["list"]]]], 1)
      [:hello, :team, "this", "is", "now", ["flat", ["list"]]]

      iex> flatten_levels([[:hello, :team], ["this", "is"], ["now", ["flat", ["list"]]]], 2)
      [:hello, :team, "this", "is", "now", "flat", ["list"]]

      iex> flatten_levels([[:hello, :team], ["this", "is"], ["now", ["flat", ["list"]]]], 3)
      [:hello, :team, "this", "is", "now", "flat", "list"]

      iex> flatten_levels([[:hello, :team], ["this", "is"], ["now", ["flat", ["list"]]]], 4)
      [:hello, :team, "this", "is", "now", "flat", "list"]

      iex> flatten_levels([[:hello, :team], ["this", "is"], ["now", ["flat", ["list"]]]], 0)
      [[:hello, :team], ["this", "is"], ["now", ["flat", ["list"]]]]
  """
  def flatten_levels(list, levels) do
    list
    |> mark_list(levels)
    |> List.flatten()
    |> unmark_list()
  end

  defp mark_list(list, 0) do
    Enum.map(list, fn item -> {:__flatten_mark, item} end)
  end

  defp mark_list(list, level) do
    Enum.map(list, fn
      list when is_list(list) -> mark_list(list, level - 1)
      other -> other
    end)
  end

  defp unmark_list(list) do
    Enum.map(list, fn
      {:__flatten_mark, item} -> item
      list when is_list(list) -> unmark_list(list)
      other -> other
    end)
  end

  defp each_if_ok_reducer(item, :ok, fun) do
    with :ok <- fun.(item) do
      {:ok, :ok}
    end
  end

  defp halt_on_error(item, {:ok, acc}, fun) do
    case fun.(item, acc) do
      {:ok, _acc} = new -> {:cont, new}
      {:error, _reason} = error -> {:halt, error}
    end
  end

  defp map_if_ok_reducer(item, acc, fun) do
    with {:ok, value} <- fun.(item) do
      {:ok, [value | acc]}
    end
  end
end