union.ex

defmodule MyApp.Types.Union do
  @constraints [
    types: [
      type: {:custom, __MODULE__, :union_types, []},
      doc: """
      The types to be unioned, a map of a string name for the enum value to its configuration.

      For example:

          %{
            "int" => %{
              type: :integer,
              constraints: [
                max: 10
              ]
            },
            "object" => %{
              type: MyObjectType,
              tag: :type,
              tag_value: "my_object"
            },
            "other_object" => %{
              type: MyOtherObjectType,
              tag: :type,
              tag_value: "my_other_object"
            },
            "other_object_without_type"=> %{
               type: MyOtherObjectTypeWithoutType,
              tag: :type,
              tag_value: nil
            }
          }
      """
    ]
  ]

  @doc false
  def union_types(value) do
    {:ok,
     Enum.reduce(value, %{}, fn {name, config}, types ->
       config =
         Map.update!(config, :type, fn type ->
           Ash.Type.get_type(type)
         end)

       if config[:type] == __MODULE__ do
         config
         |> Map.update!(:constraints, fn constraints ->
           Keyword.update!(constraints, :types, fn types ->
             {:ok, types} = union_types(types)
             types
           end)
         end)
         |> Map.get(:constraints)
         |> Keyword.get(:types)
         |> Enum.reduce(types, fn {new_name, new_config}, types ->
           if types[new_name] do
             raise "Detected a conflict in nested union type names. They must be unique all the way down."
           else
             Map.put(types, new_name, new_config)
           end
         end)
       else
         Map.put(types, name, config)
       end
     end)
     |> Map.new(fn {key, config} ->
       type = Ash.Type.get_type(config[:type])

       case type do
         {:array, type} ->
           case Code.ensure_compiled(type) do
             {:module, _} ->
               if !Ash.Type.ash_type?(type) do
                 raise """
                 Unknown type in union type \"#{key}\": #{inspect(config[:type])}
                 """
               end

             _ ->
               nil
           end

         type ->
           case Code.ensure_compiled(type) do
             {:module, _} ->
               if !Ash.Type.ash_type?(type) do
                 raise """
                 Unknown type in union type \"#{key}\": #{inspect(config[:type])}
                 """
               end

             _ ->
               nil
           end
       end

       schema = Ash.Type.constraints(type)
       constraints = Ash.OptionsHelpers.validate!(config[:constraints] || [], schema)

       {key, Map.put(config, :constraints, constraints)}
     end)}
  end

  @moduledoc """
  A union between multiple types, distinguished with a tag or by attempting to validate.

  ## Constraints

  #{Ash.OptionsHelpers.docs(@constraints)}
  """
  use Ash.Type

  @type t :: %__MODULE__{}
  defstruct [:value, :type]

  defimpl Jason.Encoder do
    def encode(%{value: value}, opts) do
      Jason.Encode.value(value, opts)
    end
  end

  @impl true
  def constraints, do: @constraints

  ## Find the minimal supported type?
  @impl true
  def storage_type, do: :map

  # @impl true
  # def prepare_change(%__MODULE__{value: old_value, type: type_name}, new_value, constraints) do
  #   if is_map(old_value) && is_map(new_value) && tag = constraints[:types][type_name][:tag] do
  #     if !Map.has_key?(new_value, tag) && !Map.has_key?(new_value, to_string(tag)) do

  #     end
  #   end
  # end

  @impl true
  def cast_input(nil, _), do: {:ok, nil}

  def cast_input(%__MODULE__{value: value, type: type_name}, constraints) do
    type = constraints[:types][type_name][:type]
    inner_constraints = constraints[:types][type_name][:constraints] || []

    case Ash.Type.cast_input(type, value, inner_constraints) do
      {:ok, value} ->
        {:ok, %__MODULE__{value: value, type: type_name}}

      error ->
        error
    end
  end

  def cast_input(value, constraints) do
    types = constraints[:types] || %{}

    types
    |> Enum.sort_by(fn {_type_name, config} -> config[:tag] end)
    |> Enum.reverse()
    |> Enum.reduce_while({:error, %{}}, fn {type_name, config}, {:error, errors} ->
      type = config[:type]

      if is_map(value) && config[:tag] do
        tag_value = config[:tag_value]

        value =
          if Ash.Type.embedded_type?(type) && !is_struct(value) do
            Enum.reduce(Ash.Resource.Info.attributes(type), value, fn attr, value ->
              with {:array, _nested_type} <- attr.type,
                   true <- has_key?(value, attr.name) do
                update_key(value, attr.name, fn value ->
                  if is_map(value) && !is_struct(value) do
                    Map.values(value)
                  else
                    value
                  end
                end)
              else
                _ ->
                  value
              end
            end)
          else
            value
          end

        if (Map.get(value, config[:tag]) || Map.get(value, to_string(config[:tag]))) == tag_value do
          case Ash.Type.cast_input(type, value, config[:constraints] || []) do
            {:ok, value} ->
              {:halt,
               {:ok,
                %__MODULE__{
                  value: value,
                  type: type_name
                }}}

            {:error, other} ->
              {:halt, {:error, "is not a valid #{type_name}: #{inspect(other)}"}}

            :error ->
              {:halt, {:error, "is not a valid #{type_name}"}}
          end
        else
          {:cont,
           {:error,
            Map.put(errors, type_name, "#{config[:tag]} does not equal #{config[:tag_value]}")}}
        end
      else
        if config[:tag] do
          {:cont, {:error, Map.put(errors, type_name, "is not a map")}}
        else
          case Ash.Type.cast_input(type, value, config[:constraints] || []) do
            {:ok, value} ->
              {:halt,
               {:ok,
                %__MODULE__{
                  value: value,
                  type: type_name
                }}}

            {:error, other} ->
              {:cont, {:error, Map.put(errors, type_name, other)}}

            :error ->
              {:cont, {:error, Map.put(errors, type_name, "is invalid")}}
          end
        end
      end
    end)
    |> case do
      {:error, errors} when is_binary(errors) ->
        {:error, errors}

      {:error, errors} ->
        {:error, error_message(errors)}

      {:ok, value} ->
        {:ok, value}

      value ->
        value
    end
  end

  defp has_key?(map, key) do
    Map.has_key?(map, key) || Map.has_key?(map, to_string(key))
  end

  defp update_key(map, key, func) do
    cond do
      Map.has_key?(map, key) ->
        Map.update!(map, key, func)

      Map.has_key?(map, to_string(key)) ->
        Map.update!(map, to_string(key), func)

      true ->
        map
    end
  end

  defp error_message(errors) do
    "No union type matched\n" <>
      Enum.map_join(errors, "\n", fn {key, errors} ->
        "  #{key}: #{inspect(errors)}"
      end)
  end

  @impl true
  def cast_stored(nil, _), do: {:ok, nil}

  def cast_stored(%{"type" => type, "value" => value}, constraints) do
    types = constraints[:types] || %{}

    case Map.fetch(types, type) do
      {:ok, config} ->
        case Ash.Type.cast_stored(config[:type], value, config[:constraints]) do
          {:ok, casted_value} ->
            {:ok,
             %__MODULE__{
               value: casted_value,
               type: type
             }}

          other ->
            other
        end

      other ->
        other
    end
  end

  def cast_stored(_, _), do: :error

  @impl true
  def dump_to_native(nil, _), do: {:ok, nil}

  def dump_to_native(%__MODULE__{value: value, type: type_name}, constraints) do
    type = constraints[:types][type_name][:type]
    constraints = constraints[:types][type_name][:constraints] || []

    if type do
      case Ash.Type.dump_to_native(type, value, constraints) do
        {:ok, value} ->
          {:ok, %{"type" => type_name, "value" => value}}

        other ->
          other
      end
    else
      :error
    end
  end
end