defmodule PropEx.TypespecClause do
  defmacro typespec_clause(name, args \\ []) do
    fun_args = Enum.map(args, fn name -> {:var!, [], [{name, [], __MODULE__}]} end)
    call_args = Enum.map(args, fn name -> {:convert_typespec, [], [{:var!, [], [{name, [], __MODULE__}]}]} end)
    call = {{:., [], [:proper_types, name]}, [], call_args}
    quote do
      def convert_typespec({unquote(name), _, unquote(fun_args)}), do: unquote(call)
    end
  end
end

defmodule PropEx do
  @moduledoc """
  A wrapper around PropEr
  """

  alias Code.Typespec
  alias :proper_types, as: PT
  import __MODULE__.TypespecClause

  def convert_typespec(typespec)

  # straightforward types
  typespec_clause :any
  typespec_clause :term
  typespec_clause :arity
  typespec_clause :binary
  typespec_clause :integer
  typespec_clause :neg_integer
  typespec_clause :non_neg_integer
  typespec_clause :float
  typespec_clause :list
  typespec_clause :list, [:elem]

  # type defined in another module
  def convert_typespec({{:., _, [mod, type_name]}, _, _args} = ast) do
    with {:ok, types} <- Typespec.fetch_types(mod),
                types <- Enum.map(types, fn {:type, type} -> type end),
                type <- List.keyfind(types, type_name, 0),
                {:"::", _, [{^type_name, _, type_args}, type_definition]} <- Typespec.type_to_quoted(type) do
      if length(type_args) != 0, do:
        raise ArgumentError, "PropEx: Sorry, not implemented: type #{Macro.to_string(ast)} is generic"
      convert_typespec(type_definition)
    else
      error -> raise ArgumentError, "PropEx: unable to access @type #{Macro.to_string(ast)}: #{error}"
    end
  end

  # unimplemented types
  def convert_typespec({type, _, []}) when type == :pid or type == :port or type == :reference, do:
    raise ArgumentError, "PropEx: type #{type}() cannot be generated by PropEr"
  def convert_typespec(type), do:
    raise ArgumentError, "PropEx: Sorry, not implemented: type #{Macro.to_string(type)} is not supported\nRepresentation: #{inspect type}"

  def convert_argument_spec({:"::", _, [{name, _, nil}, type]}), do:
    {name, convert_typespec(type)}
  def convert_argument_spec({:"::", _, [name, type]}) when is_atom(name), do:
    {name, convert_typespec(type)}
  def convert_argument_spec(_spec), do:
    raise ArgumentError, """
    PropEx: Your type specification seems to be missing variable names.
    When using arguments_of, instead of:

        @spec add(integer(), integer()) :: integer()
        def add(a, b), do: a + b

    You should write:

        @spec add(a :: integer(), b :: integer()) :: integer()
        def add(a, b), do: a + b
    """

  def convert_argument_list(argument_list) do
    argument_list
      |> Enum.map(&convert_argument_spec/1)
      |> Enum.into(%{})
      |> Map.values
      |> :erlang.list_to_tuple
  end

  def get_arguments_of({:/, _, [call, arity]} = ast, context) do
    # parse the function reference
    {orig_mod, mod, fun, arity} = case Macro.decompose_call(call) do
      {mod, fun, []} -> {mod, Macro.expand(mod, context), fun, arity}
      _ -> get_arguments_of(:yolo, context) # to raise an error
    end

    # fetch function specification and extract the arguments from it
    args = with {:ok, specs} <- Typespec.fetch_specs(mod),
         {{^fun, ^arity}, [spec]} <- List.keyfind(specs, {fun, arity}, 0),
         {:"::", _, [{^fun, _, args}, _return_type]} <- Typespec.spec_to_quoted(fun, spec) do
      args
    else
      error -> raise ArgumentError, "PropEx: unable to access the @spec for function #{Macro.to_string(ast)}: #{error}"
    end

    # generate the AST for `arg1, arg2, ...`
    call_args = args
      |> Enum.map(fn
        {:"::", _, [{name, _, nil}, _type]} -> name
        {:"::", _, [name, _type]} -> name
      end)
      |> Enum.map(fn var_name -> {:var!, [], [{var_name, [], __MODULE__}]} end)

    # generate the AST for `result = mod.fun(args)`
    result_fetcher = {:=, [], [{:var!, [], [{:result, [], __MODULE__}]}, {{:., [], [orig_mod, fun]}, [], call_args}]}

    {args, result_fetcher}
  end

  def get_arguments_of(_, _), do: raise ArgumentError, """
    PropEx: It looks like you passed something other than Mod.fun/arity to arguments_of. You should use it like this:

        forall arguments_of MyModule.my_fun/2 do
          # code
        end
    """

  defmacro __using__(_env) do
    quote do
      require PropEx
      import PropEx, only: [forall: 2]
    end
  end

  defmacro forall(argument_list, opts) do
    # if argument_list is actually a request to fetch the arguments from the
    # specs, do that
    {argument_list, result_fetcher} = case argument_list do
      {:arguments_of, _, [ast_fun_ref]} -> get_arguments_of(ast_fun_ref, __CALLER__)
      _ -> {argument_list, nil}
    end

    # process argument list
    destructurer = argument_list
      |> Enum.map(&convert_argument_spec/1)
      |> Enum.into(%{})
      |> Map.keys
      |> Enum.map(fn var_name -> {:var!, [], [{var_name, [], __MODULE__}]} end)
      |> :erlang.list_to_tuple

    quote do
      assert :proper.forall(PropEx.convert_argument_list(unquote(Macro.escape(argument_list))), fn arguments ->
        unquote(destructurer) = arguments
        unquote(result_fetcher)
        unquote(opts[:do])
      end) |> :proper.quickcheck
    end
  end
end