%-*-Mode:erlang;coding:utf-8;tab-width:4;c-basic-offset:4;indent-tabs-mode:()-*-
%%! -config /etc/cloudi/vm.config -nocookie
% ex: set ft=erlang fenc=utf-8 sts=4 ts=4 sw=4 et nomod:

-mode(compile).

-define(CMD_PING,       "ping").
-define(CMD_TEST,       "test").
-define(CMD_STOP,       "stop").
-define(CMD_RESTART,    "restart").
-define(CMD_REBOOT,     "reboot").
-define(CMD_ARGS_REMSH, "args_remsh").

-define(SHUTDOWN_TIME_MAX, 65000).

% twice this number of atoms is created when running this script
% (used for nodetool script node names and remsh node names)
-define(MAX_ATOMS, 251). % prime number

% set during configuration
-define(CONFIG_ROOT_DIR, "/usr/lib/cloudi-2.0.7").
-define(CONFIG_EPMD_PATH, ?CONFIG_ROOT_DIR "/erts-15.2.7.7/bin/epmd").
-define(CONFIG_VMARGS_PATH, ?CONFIG_ROOT_DIR "/etc/vm.args").
-define(CONFIG_VMARGS_PARSE, 'exact').

-record(state,
        {
            timeout = ?SHUTDOWN_TIME_MAX :: pos_integer(),
            test :: boolean(),
            stderr :: port(),
            heart = false :: boolean(),
            node_name_type = undefined :: longnames | shortnames | undefined,
            node = undefined :: node() | undefined,
            node_script = undefined :: node() | undefined,
            cookie = undefined :: atom()
        }).

main(Args) ->
    State0 = #state{test = test_mode(Args),
                    stderr = erlang:open_port({fd, 0, 2}, [out, {line, 256}])},
    ok = ensure_epmd_started(?CONFIG_EPMD_PATH, State0),
    StateN = initialize(parse_vmargs(?CONFIG_VMARGS_PATH, State0)),
    ok = ensure_node_connection_works(StateN),
    ok = command_line(Args, StateN),
    exit_code(0).

initialize(#state{node_name_type = NodeNameType,
                  node = Node,
                  cookie = Cookie} = State) ->
    NodeScript = node_script(Node),
    case net_kernel:start([NodeScript, NodeNameType]) of
        {ok, _} ->
            true = erlang:set_cookie(node(), Cookie),
            NodeCanonical = node_canonical(Node),
            State#state{node = NodeCanonical,
                        node_script = NodeScript};
        {error, Reason} ->
            exit_error(1, "net_kernel failed (~p)!\n", [Reason], State)
    end.

ensure_epmd_started(EpmdPath, State) ->
    case shell([EpmdPath, " -names"]) of
        {0, _} ->
            ok;
        {ExitCode, Output} ->
            exit_error(ExitCode, Output, State)
    end.

ensure_node_connection_works(#state{node = Node} = State) ->
    case net_kernel:hidden_connect_node(Node) of
        true ->
            case net_adm:ping(Node) of
                pong ->
                    ok;
                pang ->
                    exit_error(1, "Node ping failed!\n", State)
            end;
        false ->
            exit_error(1, "Node connect failed!\n", State)
    end.

command_line([?CMD_PING], _) ->
    io:format("pong\n");
command_line([?CMD_TEST], _) ->
    ok;
command_line([?CMD_STOP], State) ->
    rpc_call_init(stop, State);
command_line([?CMD_RESTART], State) ->
    rpc_call_init(restart, State);
command_line([?CMD_REBOOT],
             #state{heart = Heart} = State) ->
    if
        Heart =:= true ->
            rpc_call_init(reboot, State);
        Heart =:= false ->
            exit_error(1,
                       "reboot requires start with "
                       "-heart (in vm.args)\n", State)
    end;
command_line([?CMD_ARGS_REMSH],
             #state{node_name_type = NodeNameType,
                    node = Node,
                    node_script = NodeScript,
                    cookie = Cookie}) ->
    NameArg = if
        NodeNameType =:= shortnames ->
            "-sname";
        NodeNameType =:= longnames ->
            "-name"
    end,
    NodeRemsh = node_remsh(NodeScript),
    io:format("~s ~s -remsh ~s -setcookie ~s\n",
              [NameArg, NodeRemsh, Node, shell_atom(Cookie)]);
command_line(_, State) ->
    exit_error(1, help(), State).

node_canonical(Node) ->
    case node_split(Node) of
        {Name, ""} ->
            {_, Host} = node_split(node()),
            erlang:list_to_atom(Name ++ [$@ | Host]);
        {_, _} ->
            Node
    end.

node_script(Node) ->
    {Name, Host} = node_split(Node),
    NameUnique = Name ++ "_script_process" ++ unique(),
    if
        Host == [] ->
            erlang:list_to_atom(NameUnique);
        true ->
            erlang:list_to_atom(NameUnique ++ [$@ | Host])
    end.

node_remsh(NodeScript) ->
    {NameScript, Host} = node_split(NodeScript),
    NameUniqueRemsh = NameScript ++ "remsh",
    if
        Host == [] ->
            erlang:list_to_atom(NameUniqueRemsh);
        true ->
            erlang:list_to_atom(NameUniqueRemsh ++ [$@ | Host])
    end.

unique() ->
    erlang:integer_to_list(erlang:list_to_integer(os:getpid()) rem ?MAX_ATOMS).

parse_vmargs(VMArgs, State) ->
    case file:open(VMArgs, [read, raw, read_ahead]) of
        {ok, F} ->
            parse_vmargs_line(F, State);
        {error, Reason} ->
            exit_error(1, "vm.args error (~p)!\n", [Reason], State)
    end.

parse_vmargs_line(F, State) ->
    case file:read_line(F) of
        eof ->
            ok = file:close(F),
            State;
        {ok, "-heart" ++ _} ->
            parse_vmargs_line(F, State#state{heart = true});
        {ok, "-sname" ++ SName} ->
            SNameValue = vmarg_to_atom(SName),
            parse_vmargs_line(F, State#state{node_name_type = shortnames,
                                             node = SNameValue});
        {ok, "-name" ++ Name} ->
            NameValue = vmarg_to_atom(Name),
            parse_vmargs_line(F, State#state{node_name_type = longnames,
                                             node = NameValue});
        {ok, "-setcookie" ++ SetCookie} ->
            CookieValue = vmarg_to_atom(SetCookie),
            parse_vmargs_line(F, State#state{cookie = CookieValue});
        {ok, _} ->
            parse_vmargs_line(F, State);
        {error, Reason} ->
            exit_error(1, "vm.args invalid (~p)!\n", [Reason], State)
    end.

rpc_call_init(F, #state{timeout = Timeout,
                        node = Node} = State) ->
    case rpc:call(Node, init, F, [], Timeout) of
        ok ->
            ok;
        Error ->
            exit_error(1, "~p\n~s failed!\n", [Error, F], State)
    end.

-spec shell(Exec :: iodata()) ->
    {non_neg_integer(), list(binary())}.

shell(Exec) ->
    Shell = erlang:open_port({spawn_executable, "/bin/sh"},
                             [{args, ["-"]}, {cd, "/"},
                              stream, binary, stderr_to_stdout, exit_status]),
    true = erlang:port_command(Shell, ["exec ", Exec, "\n"]),
    shell_output(Shell, []).

shell_output(Shell, Output) ->
    receive
        {Shell, {data, Data}} ->
            shell_output(Shell, [Data | Output]);
        {Shell, {exit_status, Status}} ->
            {Status, lists:reverse(Output)}
    end.

vmarg_to_atom(String) ->
    erlang:list_to_atom(unquote(vmarg(String))).

unquote([H | T] = L) when H == $'; H == $" ->
    case lists:reverse(T) of
        [H | Unquoted] ->
            lists:reverse(Unquoted);
        _ ->
            L
    end;
unquote(L) ->
    L.

vmarg(L0) ->
    [H0 | T0] = vmarg_left(L0),
    case vmarg_left(lists:reverse(T0)) of
        [H1 | T1] ->
            [H0 | lists:reverse([H1 | vmarg_middle(?CONFIG_VMARGS_PARSE, T1)])];
        L1 ->
            [H0 | L1]
    end.

vmarg_left([H | T]) when H == $\s; H == $\t; H == $\r; H == $\n ->
    vmarg_left(T);
vmarg_left(L) ->
    L.

-spec vmarg_middle(exact | space_merge, L :: string()) -> string().

vmarg_middle(exact, L) ->
    % Erlang/OTP >= 23.0
    L;
vmarg_middle(space_merge, L) ->
    % Erlang/OTP < 23.0 (https://bugs.erlang.org/browse/ERL-1051)
    vmarg_middle_space_merge(L).

vmarg_middle_space_merge([$\t | T]) ->
    vmarg_middle_space_merge([$\s | T]);
vmarg_middle_space_merge([$\s, H | T]) when H == $\s; H == $\t ->
    vmarg_middle_space_merge([$\s | T]);
vmarg_middle_space_merge([H | T]) ->
    [H | vmarg_middle_space_merge(T)];
vmarg_middle_space_merge(L) ->
    L.

shell_atom(A) ->
    shell_atom_escape(erlang:atom_to_list(A)).

shell_atom_escape([] = L) ->
    L;
shell_atom_escape([H | T])
    when (H >= $a) andalso (H =< $z);
         (H >= $A) andalso (H =< $Z);
         (H >= $0) andalso (H =< $9);
         H == $,; H == $.; H == $_; H == $+; H == $:;
         H == $@; H == $%; H == $/; H == $- ->
    [H | shell_atom_escape(T)];
shell_atom_escape([H | T]) ->
    [$\\, H | shell_atom_escape(T)].

node_split(Node) when is_atom(Node) ->
    node_split(erlang:atom_to_list(Node), []).

node_split([], Name) ->
    {lists:reverse(Name), []};
node_split([$@ | NodeStr], Name) ->
    {lists:reverse(Name), NodeStr};
node_split([C | NodeStr], Name) ->
    node_split(NodeStr, [C | Name]).

test_mode(Args) ->
    (Args == [?CMD_TEST]) orelse (Args == [?CMD_ARGS_REMSH]).

exit_error(ExitCode, Format, State) ->
    exit_error(ExitCode, Format, undefined, State).

exit_error(_, _, _,
           #state{test = true}) ->
    exit_code(1);
exit_error(ExitCode, Format, Args,
           #state{test = false,
                  stderr = STDERR}) ->
    if
        Args =:= undefined ->
            erlang:port_command(STDERR, Format);
        true ->
            erlang:port_command(STDERR, io_lib:format(Format, Args))
    end,
    exit_code(ExitCode).

exit_code(ExitCode) when is_integer(ExitCode), ExitCode >= 0 ->
    erlang:halt(ExitCode).

help() ->
    "Usage: nodetool "
        ?CMD_PING "|"
        ?CMD_TEST "|"
        ?CMD_STOP "|"
        ?CMD_RESTART "|"
        ?CMD_REBOOT "|"
        ?CMD_ARGS_REMSH
        "\n".

