-define(TIMEOUT, 60000).

-behavior(ssh_channel).

%% API
-export([start_link/2, send/2, send_wait/2, close/1, get_data/1]).

%% ssh_channel API
-export([init/1, handle_call/3, handle_cast/2, handle_msg/2]).
-export([handle_ssh_msg/2, code_change/3, terminate/2]).

-record(state, {cn,
                chan,
                shell_started=false,
                caller,
                buf=[]}).

start_link(ConnInfo, Host) ->
  User = proplists:get_value(user, ConnInfo),
  Password = proplists:get_value(password, ConnInfo),
  application:start(crypto),
  application:start(ssh),
  case ssh:connect(Host, 22, [{user, User},
                              {password, Password},
                              {silently_accept_hosts, true}]) of
    {ok, Cn} ->
      case ssh_connection:session_channel(Cn, ?TIMEOUT) of
        {ok, ChanId} ->
          ssh_channel:start_link(Cn, ChanId, ?MODULE, [Cn, ChanId]);
        Error ->
          ssh:close(Cn),
          throw(Error)
      end;
    Error ->
      throw(Error)
  end.

get_data(ChannelPid) ->
  ssh_channel:call(ChannelPid, get_data, infinity).

send(ChannelPid, Data) ->
  ssh_channel:call(ChannelPid, {send, Data}, infinity).

send_wait(ChannelPid, Data) ->
  Me = self(),
  ssh_channel:call(ChannelPid, {send_wait, Me, Data}, infinity),
  receive
    {response, Response} ->
      Response
  after ?TIMEOUT ->
      {error, timeout}
  end.

close(ChannelPid) ->
  ssh_channel:cast(ChannelPid, close).

init([Conn, ChanId]) ->
  {ok, #state{cn=Conn, chan=ChanId}}.

handle_call(get_data, _From, #state{buf=Buf}=State) ->
  {reply, Buf, State#state{buf=[]}};

handle_call({send, Data}, _From, #state{cn=Cn, chan=Chan}=State) ->
  {reply, ssh_connection:send(Cn, Chan, Data, ?TIMEOUT), State};

handle_call({send_wait, Caller, Data}, _From, #state{cn=Cn, chan=Chan, shell_started=ShellStarted}=State) ->
  case  ShellStarted =:= true orelse ssh_connection:shell(Cn, Chan) =:= ok of
    true ->
      NewState = State#state{shell_started=true},
      case ssh_connection:send(Cn, Chan, Data, ?TIMEOUT) of
        ok ->
          {reply, ok, NewState#state{caller=Caller}};
        Error ->
          {reply, Error, NewState}
      end;
    false ->
      {reply, {error, no_shell}, State}
  end;

handle_call(_Msg, _From, State) ->
  {reply, ignored, State}.

handle_cast(close, State) ->
  {stop, normal, State};

handle_cast(_Msg, State) ->
  {noreply, State}.

handle_msg(_Msg, State) ->
  io:format("Msg: ~p~n", [_Msg]),
  {ok, State}.

handle_ssh_msg({ssh_cm, _, {data, _, _, Data}}, #state{caller=Caller}=State) when not(Caller =:= undefined) ->
  Caller ! {response, Data},
  {ok, State#state{caller=undefined}};

handle_ssh_msg({ssh_cm, _, {data, _, _, Data}}, #state{buf=Buf}=State) ->
  {ok, State#state{buf=[Data|Buf]}};

handle_ssh_msg(_Msg, State) ->
  io:format("SSH Message: ~p~n", [_Msg]),
  {ok, State}.

code_change(_OldVsn, State, _Extra) ->
  {ok, State}.

terminate(_Reason, _State) ->
  ok.