ecto_postgres_enum

defmodule EctoPostgresEnum do
  @moduledoc """
  Helper module to define enum for `ecto` and `PostgreSQL` with support for dynamic values.

  ## Usage

      defmodule MyEnum do
        values = [:my, :enum]
        use EctoPostgresEnum, schema: :my_schema, type: :my_type, values: values
      end

  ## Options

  * `schema` - Allows to change [PostgreSQL Schema](https://www.postgresql.org/docs/current/ddl-schemas.html) for specified enum.
  * `type` - Allows to change type identifier. It's useful for migrations and required by `c:Ecto.Type.type/0` callback.

  Name   | Type       | Required | Default
  :----- | :--------- | :------- | :------------------------------------------------
  schema | atom       | false    | nil (fallbacks to `PostgreSQL` default: "public")
  type   | atom       | false    | atom (underscored last module part)
  values | list(atom) | true     | N/A

  ## Debug
  This library automatically generates few useful functions to work with allowed values.

  Example usage:

      defmodule MyEnum do
        values = [:my, :enum]
        use EctoPostgresEnum, schema: :my_schema, type: :my_type, values: values
      end

      defmodule MyApp do
        require MyEnum

        def handle_user_input(input) when MyEnum.valid_string?(input) do
          IO.puts "Alright, \#{input} is allowed value!"
        end

        def handle_user_input(input) do
          IO.puts "Ooops, \#{input} is not allowed value! Please try again …"
        end
      end

      MyApp.handle_user_input("my")
      {:ok, "my"}

      MyApp.handle_user_input("something")
      {:error. "Wrong enum value!"}
  """

  @doc false
  def gen_default_type(module),
    do: module |> Module.split() |> List.last() |> Macro.underscore() |> String.to_atom()

  @doc false
  def gen_type_ast(values), do: values |> Enum.reverse() |> do_gen_type_ast()

  defp do_gen_type_ast([head | tail]), do: Enum.reduce(tail, head, &{:|, [], [&1, &2]})

  @doc false
  defmacro __using__(opts), do: [base(opts), database_block(), debug_block(), type()]

  defp base(opts), do: [base_definitions_block(opts), base_checks_block(), base_rest_block()]

  defp base_definitions_block(opts) do
    quote bind_quoted: [opts: opts], unquote: false do
      is_list(opts) || raise "Options should be a list!"

      @behaviour Ecto.Type
      @__schema__ opts[:schema]
      @__type__ opts[:type] || EctoPostgresEnum.gen_default_type(__MODULE__)
      @__values__ opts[:values]
    end
  end

  defp base_checks_block do
    quote do
      @__values__ || raise "Option values (list) is required!"
      is_atom(@__type__) || raise "Type needs to be an atom!"
      @__values__ == Enum.uniq(@__values__) || raise "Duplicates are not allowed in enum values!"
      Enum.count(@__values__) > 0 || raise "Valid enums requires at least 1 different values!"
      Enum.all?(@__values__, &is_atom/1) || raise "All values must be atoms!"
      @__schema__ && (is_atom(@__schema__) || raise "Option schema must be atom!")
    end
  end

  defp base_rest_block do
    quote unquote: false do
      alias Ecto.Migration

      type_sql = if is_nil(@__schema__), do: @__type__, else: :"#{@__schema__}.#{@__type__}"
      values_sql = Enum.map_join(@__values__, ", ", &"'#{&1}'")
      @__create_sql__ "CREATE TYPE #{type_sql} AS ENUM (#{values_sql})"
      @__drop_sql__ "DROP TYPE #{type_sql}"
      @__string_values__ Enum.map(@__values__, &Atom.to_string/1)
      @__zipped_values__ Enum.zip(@__values__, @__string_values__)

      @type t :: unquote(EctoPostgresEnum.gen_type_ast(@__values__))
    end
  end

  defp database_block do
    quote do
      @doc "Creates database enum."
      @spec create_db_enum :: :ok
      def create_db_enum, do: Migration.execute(@__create_sql__)

      @doc "Drops database enum."
      @spec drop_db_enum :: :ok
      def drop_db_enum, do: Migration.execute(@__drop_sql__)
    end
  end

  defp debug_block do
    quote do
      @doc "Returns list of allowed atom values."
      @spec atom_values :: [t]
      def atom_values, do: @__values__

      @doc "Returns list of allowed string values."
      @spec string_values :: [String.t()]
      def string_values, do: @__string_values__

      @doc "Checks if given atom is allowed."
      @spec valid_atom?(t) :: true
      @spec valid_atom?(term) :: false
      defguard valid_atom?(value) when value in @__values__

      @doc "Checks if given string is allowed."
      @spec valid_string?(String.t()) :: boolean
      @spec valid?(term) :: false
      defguard valid_string?(value) when value in @__string_values__

      @doc "Checks if given atom or string is allowed."
      @spec valid?(t) :: true
      @spec valid?(String.t()) :: boolean
      @spec valid?(term) :: false
      defguard valid?(value) when valid_atom?(value) or valid_string?(value)

      @doc "Returns zipped list of atom values with list of string values."
      @spec values :: [term]
      def values, do: @__zipped_values__
    end
  end

  defp type, do: [type_cast(), type_dump(), type_rest()]

  defp type_cast do
    quote unquote: false do
      @doc """
      Casts the given input to the custom type.

      This is callback implementation for: `c:Ecto.Type.cast/1`
      """
      @impl Ecto.Type
      @spec cast(String.t()) :: {:ok, t} | :error
      for {atom, string} <- @__zipped_values__ do
        @spec cast(unquote(atom)) :: {:ok, unquote(atom)}
        def cast(unquote(atom)), do: {:ok, unquote(atom)}
        def cast(unquote(string)), do: {:ok, unquote(atom)}
      end

      @spec cast(term) :: :error
      def cast(_term), do: :error
    end
  end

  defp type_dump do
    quote unquote: false do
      @doc """
      Dumps the given term into an Ecto native type.

      This is callback implementation for: `c:Ecto.Type.dump/1`
      """
      @impl Ecto.Type
      @spec dump(t) :: {:ok, String.t()}
      @spec dump(String.t()) :: {:ok, String.t()} | :error
      for {atom, string} <- @__zipped_values__ do
        def dump(unquote(atom)), do: {:ok, unquote(string)}
        def dump(unquote(string)), do: {:ok, unquote(string)}
      end

      @spec dump(term) :: :error
      def dump(_term), do: :error
    end
  end

  defp type_rest do
    quote unquote: false do
      @doc """
      Loads the given term into a custom type.

      This is callback implementation for: `c:Ecto.Type.load/1`
      """
      @impl Ecto.Type
      @spec load(String.t()) :: {:ok, t} | :error
      for {atom, string} <- @__zipped_values__ do
        def load(unquote(string)), do: {:ok, unquote(atom)}
      end

      @spec load(term) :: :error
      def load(_value), do: :error

      @doc """
      Returns the underlying schema type for the custom type.

      This is callback implementation for: `c:Ecto.Type.type/0`
      """
      @impl Ecto.Type
      @spec type :: atom
      def type, do: @__type__
    end
  end
end