prometheus_ex

defmodule Prometheus.Injector do

  def inject(callback, env, ast) do
    ast
    |> Macro.prewalk(fn (thing) ->
      case thing do
        {:def, _, _} = defun -> defun
        {:in, _, _} = arrow -> arrow #otherwise e in RuntimeError will be rewritten
        _ -> Macro.expand(thing, env)
      end
    end)
    |>  inject_(callback)
  end

  # lambda
  def inject_({:fn, fn_meta, [{:->, arrow_meta, [args, do_block]}]}, callback) do
    case args do
      [] ->
        callback.({{:., [],
                    [{:fn, fn_meta,
                       [{:->, arrow_meta,
                        [[], do_block]}]}]}, [],
                   []})
      _ ->
        names = args
        |> Enum.map(fn({name, _, _}) -> name end)
        raise Prometheus.InvalidBlockArityError, args: names
    end
  end

  # do_blocks can be simple calls or defs
  def inject_([{:do, {:__block__, [], do_blocks}}], callback) do
    do_blocks = List.flatten(do_blocks)
    if have_defs(do_blocks) do
      Enum.map(do_blocks, &inject_to_def(&1, callback))
    else
       callback.({:__block__, [], do_blocks})
    end
  end

  # just do
  def inject_([{:do, do_block}], callback) do
    inject_([{:do, {:__block__, [], [do_block]}}], callback)
  end

  # implicit try
  def inject_([{:do, _do_block} | rest] = all, callback) do
    if is_try_unwrapped(rest) do
      callback.(
        quote do
          try unquote(all)
        end)
    else
      raise "Unexpected do block #{inspect rest}"
    end
  end

  # single do, or other non-block stuff like function calls
  def inject_(thing, callback) do
    inject_([{:do, {:__block__, [], [thing]}}], callback)
  end

  defp is_try_unwrapped(block) do
    Keyword.has_key?(block, :catch)
    || Keyword.has_key?(block, :rescue)
    || Keyword.has_key?(block, :after)
    || Keyword.has_key?(block, :else)
  end

  defp have_defs(blocks) do
    defs_count = Enum.count(blocks, fn
      ({:def, _, _}) -> true
      (_) -> false
    end)

    blocks_count = Enum.count(blocks)

    case defs_count do
      0 -> false
      ^blocks_count -> true
      _ -> raise "Mixing defs and other blocks isn't allowed"
    end
  end

  defp inject_to_def({:def, def_meta, [head, [do: body]]}, callback) do
    {:def, def_meta, [head, [do: callback.(body)]]}
  end
  defp inject_to_def({:def, def_meta, [head, [{:do, _do_block} | _rest] = all]}, callback) do
    {:def, def_meta, [head, [do: callback.(
                                   quote do
                                     try unquote(all)
                                   end)]]}
  end
end