#!/usr/bin/env escript
%% -*- erlang -*-

-record(state, {run_hooks = #{},
		run_fold_hooks = #{},
		hooked_funs = {#{}, #{}},
		iq_handlers = {#{}, #{}},
		exports = #{},
		module :: module(),
		file :: filename:filename()}).

main(Paths) ->
    State =
	fold_beams(
	  fun(File0, Tree, X, Acc0) ->
		  BareName = filename:rootname(filename:basename(File0)),
		  Mod = list_to_atom(BareName),
		  File = BareName ++ ".erl",
		  Exports = maps:put(Mod, X, Acc0#state.exports),
		  Acc1 = Acc0#state{file = File, module = Mod, exports = Exports},
		  erl_syntax_lib:fold(
		    fun(Form, Acc) ->
		  	    case erl_syntax:type(Form) of
		  		application ->
		  		    case erl_syntax_lib:analyze_application(Form) of
		  			{ejabberd_hooks, {run, N}} when N == 2; N == 3 ->
		  			    collect_run_hook(Form, Acc);
		  			{ejabberd_hooks, {run_fold, N}} when N == 3; N == 4 ->
		  			    collect_run_fold_hook(Form, Acc);
		  			{ejabberd_hooks, {add, N}} when N == 4; N == 5 ->
		  			    collect_run_fun(Form, add, Acc);
		  			{ejabberd_hooks, {delete, N}} when N == 4; N == 5 ->
		  			    collect_run_fun(Form, delete, Acc);
		  			{gen_iq_handler, {add_iq_handler, 5}} ->
		  			    collect_iq_handler(Form, add, Acc);
		  			{gen_iq_handler, {remove_iq_handler, 3}} ->
		  			    collect_iq_handler(Form, delete, Acc);
		  			_ ->
		  			    Acc
		  		    end;
		  		_ ->
		  		    Acc
		  	    end
		    end, Acc1, Tree)
	  end, #state{}, Paths),
    check_hooks_arity(State#state.run_hooks),
    check_hooks_arity(State#state.run_fold_hooks),
    check_iq_handlers_export(State#state.iq_handlers, State#state.exports),
    analyze_iq_handlers(State#state.iq_handlers),
    analyze_hooks(State#state.hooked_funs),
    RunDeps = build_deps(State#state.run_hooks, State#state.hooked_funs),
    RunFoldDeps = build_deps(State#state.run_fold_hooks, State#state.hooked_funs),
    emit_module(RunDeps, RunFoldDeps, hooks_type_test).

collect_run_hook(Form, State) ->
    [Hook|Tail] = erl_syntax:application_arguments(Form),
    case atom_value(Hook, State) of
	undefined ->
	    State;
	HookName ->
	    Args = case Tail of
		       [_Host, Args0] -> Args0;
		       [Args0] ->
			   Args0
		   end,
	    Arity = erl_syntax:list_length(Args),
	    Hooks = maps:put({HookName, Arity},
			     {State#state.file, erl_syntax:get_pos(Hook)},
			     State#state.run_hooks),
	    State#state{run_hooks = Hooks}
    end.

collect_run_fold_hook(Form, State) ->
    [Hook|Tail] = erl_syntax:application_arguments(Form),
    case atom_value(Hook, State) of
	undefined ->
	    State;
	HookName ->
	    Args = case Tail of
		       [_Host, _Val, Args0] -> Args0;
		       [_Val, Args0] -> Args0
		   end,
	    Arity = erl_syntax:list_length(Args) + 1,
	    Hooks = maps:put({HookName, Arity},
			     {State#state.file, erl_syntax:get_pos(Form)},
			     State#state.run_fold_hooks),
	    State#state{run_fold_hooks = Hooks}
    end.

collect_run_fun(Form, Action, State) ->
    [Hook|Tail] = erl_syntax:application_arguments(Form),
    case atom_value(Hook, State) of
	undefined ->
	    State;
	HookName ->
	    {Module, Fun, Seq} = case Tail of
				     [_Host, M, F, S] ->
					 {M, F, S};
				     [M, F, S] ->
					 {M, F, S}
				 end,
	    ModName = module_name(Module, State),
	    FunName = atom_value(Fun, State),
	    SeqInt = integer_value(Seq, State),
	    if ModName /= undefined, FunName /= undefined, SeqInt /= undefined ->
		    Pos = case Action of
			      add -> 1;
			      delete -> 2
			  end,
		    Funs = maps_append(
			     HookName,
			     {ModName, FunName, SeqInt,
			      {State#state.file, erl_syntax:get_pos(Form)}},
			     element(Pos, State#state.hooked_funs)),
		    Hooked = setelement(Pos, State#state.hooked_funs, Funs),
		    State#state{hooked_funs = Hooked};
	       true ->
		    State
	    end
    end.

collect_iq_handler(Form, add, #state{iq_handlers = {Add, Del}} = State) ->
    [Component, _Host, Namespace, Module, Function] = erl_syntax:application_arguments(Form),
    Mod = module_name(Module, State),
    Fun = atom_value(Function, State),
    Comp = atom_value(Component, State),
    NS = binary_value(Namespace, State),
    if Mod /= undefined, Fun /= undefined, Comp /= undefined, NS /= undefined ->
	    Handlers = maps_append(
			 {Comp, NS},
			 {Mod, Fun,
			  {State#state.file, erl_syntax:get_pos(Form)}},
			 Add),
	    State#state{iq_handlers = {Handlers, Del}};
       true ->
	    State
    end;
collect_iq_handler(Form, delete, #state{iq_handlers = {Add, Del}} = State) ->
    [Component, _Host, Namespace] = erl_syntax:application_arguments(Form),
    Comp = atom_value(Component, State),
    NS = binary_value(Namespace, State),
    if Comp /= undefined, NS /= undefined ->
	    Handlers = maps_append(
			 {Comp, NS},
			 {State#state.file, erl_syntax:get_pos(Form)},
			 Del),
	    State#state{iq_handlers = {Add, Handlers}};
       true ->
	    State
    end.

check_hooks_arity(Hooks) ->
    maps:fold(
      fun({Hook, Arity}, _, M) ->
	      case maps:is_key(Hook, M) of
		  true ->
		      err("Error: hook ~s is called with different "
			  "number of arguments~n", [Hook]);
		  false ->
		      maps:put(Hook, Arity, M)
	      end
      end, #{}, Hooks).

check_iq_handlers_export({HookedFuns, _}, Exports) ->
    maps:map(
      fun(_, Funs) ->
	      lists:foreach(
		fun({Mod, Fun, {File, FileNo}}) ->
			case is_exported(Mod, Fun, 1, Exports) of
			    true -> ok;
			    false ->
				err("~s:~p: Error: "
				    "iq handler is registered on unexported function: "
				    "~s:~s/1~n", [File, FileNo, Mod, Fun])
			end
		end, Funs)
      end, HookedFuns).

analyze_iq_handlers({Add, Del}) ->
    maps:map(
      fun(Handler, Funs) ->
	      lists:foreach(
		fun({_, _, {File, FileNo}}) ->
			case maps:is_key(Handler, Del) of
			    true -> ok;
			    false ->
				err("~s:~p: Error: "
				    "iq handler is added but not removed~n",
				    [File, FileNo])
			end
		end, Funs)
      end, Add),
    maps:map(
      fun(Handler, Meta) ->
	      lists:foreach(
		fun({File, FileNo}) ->
			case maps:is_key(Handler, Add) of
			    true -> ok;
			    false ->
				err("~s:~p: Error: "
				    "iq handler is removed but not added~n",
				    [File, FileNo])
			end
		end, Meta)
      end, Del).

analyze_hooks({Add, Del}) ->
    Del1 = maps:fold(
	     fun(Hook, Funs, D) ->
		     lists:foldl(
		       fun({Mod, Fun, Seq, {File, FileNo}}, D1) ->
			       maps:put({Hook, Mod, Fun, Seq}, {File, FileNo}, D1)
		       end, D, Funs)
	     end, #{}, Del),
    Add1 = maps:fold(
	     fun(Hook, Funs, D) ->
		     lists:foldl(
		       fun({Mod, Fun, Seq, {File, FileNo}}, D1) ->
			       maps:put({Hook, Mod, Fun, Seq}, {File, FileNo}, D1)
		       end, D, Funs)
	     end, #{}, Add),
    lists:foreach(
      fun({{Hook, Mod, Fun, _} = Key, {File, FileNo}}) ->
	      case maps:is_key(Key, Del1) of
		  true -> ok;
		  false ->
		      err("~s:~p: Error: "
			  "hook ~s->~s->~s is added but was never deleted~n",
			  [File, FileNo, Hook, Mod, Fun])
	      end
      end, maps:to_list(Add1)),
    lists:foreach(
      fun({{Hook, Mod, Fun, _} = Key, {File, FileNo}}) ->
	      case maps:is_key(Key, Add1) of
		  true -> ok;
		  false ->
		      err("~s:~p: Error: "
			  "hook ~s->~s->~s is deleted but was never added~n",
			  [File, FileNo, Hook, Mod, Fun])
	      end
      end, maps:to_list(Del1)).

build_deps(Hooks, {HookedFuns, _}) ->
    maps:fold(
      fun({Hook, Arity}, Meta, Deps) ->
	      case maps:find(Hook, HookedFuns) of
		  {ok, Funs} ->
		      ExportedFuns =
			  lists:map(
			    fun({M, F, Seq, FunMeta}) ->
				    {{M, F, Arity}, Seq, FunMeta}
			    end, Funs),
		      maps_append_list({Hook, Arity, Meta}, ExportedFuns, Deps);
		  error ->
		      maps_append_list({Hook, Arity, Meta}, [], Deps)
	      end
      end, #{}, Hooks).

module_name(Form, State) ->
    try
	Name = erl_syntax:macro_name(Form),
	'MODULE' = erl_syntax:variable_name(Name),
	State#state.module
    catch _:_ ->
	    atom_value(Form, State)
    end.

atom_value(Form, State) ->
    case erl_syntax:type(Form) of
	atom ->
	    erl_syntax:atom_value(Form);
	_ ->
	    warn_type(Form, State, "not an atom"),
	    undefined
    end.

integer_value(Form, State) ->
    case erl_syntax:type(Form) of
	integer ->
	    erl_syntax:integer_value(Form);
	_ ->
	    warn_type(Form, State, "not an integer"),
	    undefined
    end.

binary_value(Form, State) ->
    try erl_syntax:concrete(Form) of
	Binary when is_binary(Binary) ->
	    Binary;
	_ ->
	    warn_type(Form, State, "not a binary"),
	    undefined
    catch _:_ ->
	    warn_type(Form, State, "not a binary"),
	    undefined
    end.

is_exported(Mod, Fun, Arity, Exports) ->
    try maps:get(Mod, Exports) of
	L -> lists:member({Fun, Arity}, L)
    catch _:{badkey, _} -> false
    end.

warn_type({var, _, 'Type'}, #state{module = mod_delegation}, "not an atom") ->
    ok;
warn_type({var, _, 'NS'}, #state{module = mod_delegation}, "not a binary") ->
    ok;
warn_type({var, _, _}, #state{module = gen_mod}, _) ->
    ok;
warn_type(Form, State, Warning) ->
    log("~s:~p: Warning: " ++ Warning ++ ": ~s~n",
	[State#state.file,
	 erl_syntax:get_pos(Form),
	 erl_prettypr:format(Form)]).

emit_module(RunDeps, RunFoldDeps, Module) ->
    File = filename:join(["src", Module]) ++ ".erl",
    try
	{ok, Fd} = file:open(File, [write]),
	write(Fd,
	      "%% Generated automatically~n"
	      "%% DO NOT EDIT: run `make hooks` instead~n~n", []),
	write(Fd, "-module(~s).~n", [Module]),
	write(Fd, "-compile(nowarn_unused_vars).~n", []),
	write(Fd, "-dialyzer(no_return).~n~n", []),
	emit_export(Fd, RunDeps, "run hooks"),
	emit_export(Fd, RunFoldDeps, "run_fold hooks"),
	emit_run_hooks(Fd, RunDeps),
	emit_run_fold_hooks(Fd, RunFoldDeps),
	file:close(Fd),
	log("Module written to ~s~n", [File])
    catch _:{badmatch, {error, Reason}} ->
	    err("Error: writing to ~s failed: ~s", [File, file:format_error(Reason)])
    end.

emit_run_hooks(Fd, Deps) ->
    DepsList = lists:sort(maps:to_list(Deps)),
    lists:foreach(
      fun({{Hook, Arity, {File, LineNo}}, Funs}) ->
	      write(Fd, "%% called at ~s:~p~n", [File, LineNo]),
	      Args = string:join(
		       [[N] || N <- lists:sublist(lists:seq($A, $Z), Arity)],
		       ", "),
	      write(Fd, "~s(~s) ->~n    ", [Hook, Args]),
	      Calls = [io_lib:format("_ = ~s:~s(~s)", [Mod, Fun, Args])
		       || {{Mod, Fun, _}, _Seq, _} <- lists:keysort(2, Funs)],
	      write(Fd, "~s.~n~n",
		    [string:join(Calls ++ ["ok"], ",\n    ")])
      end, DepsList).

emit_run_fold_hooks(Fd, Deps) ->
    DepsList = lists:sort(maps:to_list(Deps)),
    lists:foreach(
      fun({{Hook, Arity, {File, LineNo}}, []}) ->
	      write(Fd, "%% called at ~s:~p~n", [File, LineNo]),
	      Args = ["Acc"|lists:duplicate(Arity - 1, "_")],
	      write(Fd, "~s(~s) -> Acc.~n~n", [Hook, string:join(Args, ", ")]);
	 ({{Hook, Arity, {File, LineNo}}, Funs}) ->
	      write(Fd, "%% called at ~s:~p~n", [File, LineNo]),
	      Args = [[N] || N <- lists:sublist(lists:seq($A, $Z), Arity - 1)],
	      write(Fd, "~s(~s) ->~n    ", [Hook, string:join(["Acc0"|Args], ", ")]),
	      {Calls, _} = lists:mapfoldl(
			     fun({{Mod, Fun, _}, _Seq, _}, N) ->
				     Args1 = ["Acc" ++ integer_to_list(N)|Args],
				     {io_lib:format("Acc~p = ~s:~s(~s)",
						    [N+1, Mod, Fun,
						     string:join(Args1, ", ")]),
				      N + 1}
			     end, 0, lists:keysort(2, Funs)),
	      write(Fd, "~s,~n", [string:join(Calls, ",\n    ")]),
	      write(Fd, "    Acc~p.~n~n", [length(Funs)])
      end, DepsList).

emit_export(Fd, Deps, Comment) ->
    DepsList = lists:sort(maps:to_list(Deps)),
    Exports = lists:map(
		fun({{Hook, Arity, _}, _}) ->
			io_lib:format("~s/~p", [Hook, Arity])
		end, DepsList),
    write(Fd, "%% ~s~n-export([~s]).~n~n",
	      [Comment, string:join(Exports, ",\n         ")]).

fold_beams(Fun, State, Paths) ->
    Paths1 = fold_paths(Paths),
    Total = length(Paths1),
    {_, State1} =
	lists:foldl(
	  fun(File, {I, Acc}) ->
		  io:format("Progress: ~B% (~B/~B)\r",
			    [round(I*100/Total), I, Total]),
		  case is_elixir_beam(File) of
		      true -> {I+1, Acc};
		      false ->
			  {AbsCode, Exports} = get_code_from_beam(File),
			  Acc2 = lists:foldl(
				   fun(Form, Acc1) ->
					   Fun(File, Form, Exports, Acc1)
				   end, Acc, AbsCode),
			  {I+1, Acc2}
		  end
	  end, {0, State}, Paths1),
    State1.

fold_paths(Paths) ->
    lists:flatmap(
      fun(Path) ->
	      case filelib:is_dir(Path) of
		  true ->
		      lists:reverse(
			filelib:fold_files(
			  Path, ".+\.beam\$", false,
			  fun(File, Acc) ->
					  [File|Acc]
			  end, []));
		  false ->
		      [Path]
	      end
      end, Paths).

is_elixir_beam(File) ->
    case filename:basename(File) of
	"Elixir" ++ _ -> true;
	_ -> false
    end.

get_code_from_beam(File) ->
    case beam_lib:chunks(File, [abstract_code, exports]) of
	{ok, {_, [{abstract_code, {raw_abstract_v1, Forms}}, {exports, X}]}} ->
	    {Forms, X};
	_ ->
	    err("No abstract code found in ~s~n", [File])
    end.

log(Format, Args) ->
    io:format(standard_io, Format, Args).

err(Format, Args) ->
    io:format(standard_error, Format, Args),
    halt(1).

write(Fd, Format, Args) ->
    file:write(Fd, io_lib:format(Format, Args)).

maps_append(K, V, M) ->
    maps_append_list(K, [V], M).

maps_append_list(K, L1, M) ->
    L2 = maps:get(K, M, []),
    maps:put(K, L2 ++ L1, M).