package click.runtime.remote;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;

import jist.runtime.JistAPI;

import org.apache.log4j.Logger;

import com.trilead.ssh2.ChannelCondition;
import com.trilead.ssh2.Connection;
import com.trilead.ssh2.Session;
import com.trilead.ssh2.StreamGobbler;

/**
 * SSH connector with password/public key authentication.
 * 
 * Properties:
 * user
 * password
 * port
 * prompt
 * keyfile
 * keyfilePass
 * timeout
 * 
 * @author kurth
 */
public class WrapSshConnector2 implements RemoteConnector {

  /** MAC logger. */
  public static final Logger log = Logger.getLogger(WrapSshConnector2.class);
  
  public static int nextRequestId = 1;

  Connection ssh;
  Session sess;
  OutputStreamWriter writer;
  InputStreamReader stdout;
  InputStreamReader stderr;

  private String user;
  private String password;
  private int port;
  private String keyfile;
  private String keyfilePass;
  private String prompt;
  private int timeout = 4000 /* ms */;
  
  private Map<Integer, AsynchRunner> mapRequests;

  public WrapSshConnector2() {
    mapRequests = new HashMap<Integer, AsynchRunner>();
  }

  /*
   * (non-Javadoc)
   * @see click.runtime.remote.RemoteConnector#init(java.util.Properties)
   */
  public void init(Properties prop) {
    user = prop.getProperty("user", null);
    password = prop.getProperty("password", null);
    port = Integer.parseInt(prop.getProperty("port", "22"));
    prompt = prop.getProperty("prompt", "root@wgt");
    keyfile = prop.getProperty("keyfile", "~/.ssh/id_rsa"); 
    keyfilePass = prop.getProperty("keyfilePass", "joespass"); // will be ignored if not needed
    timeout = Integer.parseInt(prop.getProperty("timeout", "4000"));
    
    if (null == user)
      user = System.getProperty("user.name");
    int beginIndex = keyfile.indexOf('~');
    if (beginIndex > -1)
      keyfile = System.getProperty("user.home") + keyfile.substring(beginIndex+1); 
      
  }

  /*
   * (non-Javadoc)
   * @see click.runtime.remote.RemoteConnector#connect(java.lang.String, int, java.lang.String, java.lang.String)
   */
  public void connect(String host, int port, String user, String password)
      throws IOException {
    /* Create a connection instance */

    ssh = new Connection(host);

    /* Now connect */

    ssh.connect();

    /* Authenticate.
     * If you get an IOException saying something like
     * "Authentication method password not supported by the server at this stage."
     * then please check the FAQ.
     */

    // use 3 tries, cause sometimes authentication fails for some reason...
    boolean isAuthenticated = false;
    for (int tries = 3; null != this.keyfile && tries > 0; tries--) {
      File file = new File(this.keyfile);
      if (file.exists()) {
        isAuthenticated = ssh.authenticateWithPublicKey(user, file, keyfilePass);
        if (isAuthenticated)
          break;
        log.warn("Public key authentication with key file "+this.keyfile+" failed.");
        try {
          Thread.sleep(394);
        } catch (InterruptedException e) {}
      }
    }
    

    if (!isAuthenticated && null != this.password) {
      isAuthenticated = ssh.authenticateWithPassword(user, password);
    }

    if (isAuthenticated == false)
      throw new IOException("Authentication failed.");

    /* Create a session */
    
    sess = ssh.openSession();

    int x_width = 90;
    int y_width = 30;
    sess.requestPTY("dumb", x_width, y_width, 0, 0, null);
    sess.startShell();

    writer = new OutputStreamWriter(sess.getStdin());
    stdout = new InputStreamReader(sess.getStdout());
    stderr = new InputStreamReader(sess.getStderr());

    StringBuilder builderStdOut = new StringBuilder();
    StringBuilder builderStdErr = new StringBuilder();
    waitFor(builderStdOut, builderStdErr);

    if (builderStdErr.length() != 0)
      log.warn("Error message during login:" 
          + builderStdErr.toString());
  }

  /*
   * (non-Javadoc)
   * @see click.runtime.remote.RemoteConnector#connect(java.lang.String, int)
   */
  public void connect(String host) throws IOException {
    this.connect(host, port, user, password);
  }

  /*
   * (non-Javadoc)
   * @see click.runtime.remote.RemoteConnector#disconnect()
   */
  public void disconnect() throws IOException {
    /* Close the connection */
    writer.close();
    stdout.close();
    stderr.close();
    sess.close();
    ssh.close();
  }

  /*
   * (non-Javadoc)
   * @see click.runtime.remote.RemoteConnector#execute(java.lang.String)
   */
  public String execute(String command) throws IOException {
    return execute2(command);
//    return execute1(command);
  }

  /**
   * @param command
   * @return
   * @throws IOException
   */
  private String execute1(String command) throws IOException {
    String lines = "";

    /* Create a session */
    Session sess = ssh.openSession();
    
    // TODO use input stream for command execution instead of opening sessions
    // see FAQ com.trilead
    
    try {
      sess.execCommand(command);

      InputStream stdout = new StreamGobbler(sess.getStdout());
      InputStream stderr = new StreamGobbler(sess.getStderr());

      BufferedReader stdoutReader = new BufferedReader(new InputStreamReader(stdout));
      BufferedReader stderrReader = new BufferedReader(new InputStreamReader(stderr));

      while (true)
      {
        String line = stderrReader.readLine();
        if (line == null)
          break;
        lines += line + "\n";
      }
      if (lines.length() != 0)
        log.warn("Command '"+command+"' returned an error message:" + lines);

      while (true)
      {
        String line = stdoutReader.readLine();
        if (line == null)
          break;
        lines += line + "\n";
      }

    } finally {
      /* Close this session */
      sess.close();
    }

    /* Show exit status, if available (otherwise "null") */
    Integer exitCode = sess.getExitStatus();
    if (null != exitCode && exitCode.intValue() > 0)
      throw new IOException("Command '"+command+"' returned with '"+exitCode+"'");

    return lines;
  }

  private String execute2(String command) throws IOException {
    writer.write(command + "\n");
    writer.flush();
    
    StringBuilder builderStdOut = new StringBuilder();
    StringBuilder builderStdErr = new StringBuilder();
    waitFor(builderStdOut, builderStdErr);
    
    if (builderStdErr.length() != 0)
      log.warn("Command '"+command+"' returned an error message:" 
          + builderStdErr.toString());

    return builderStdOut.toString();
  }

  /**
   * @param lines
   * @param builderStdOut
   * @param builderStdErr
   * @return
   * @throws IOException
   */
  private void waitFor(StringBuilder builderStdOut,
      StringBuilder builderStdErr) throws IOException {
    /*
     * Advanced:
     * The following is a demo on how one can read from stdout and
     * stderr without having to use two parallel worker threads (i.e.,
     * we don't use the Streamgobblers here) and at the same time not
     * risking a deadlock (due to a filled SSH2 channel window, caused
     * by the stream which you are currently NOT reading from =).
     */

    /* Don't wrap these streams and don't let other threads work on
     * these streams while you work with Session.waitForCondition()!!!
     */

    char[] buffer = new char[8192];
    

    while (true)
    {
      int avail = sess.getStdout().available();
      if (!stdout.ready() && !stderr.ready())
      {
        /* Even though currently there is no data available, it may be that new data arrives
         * and the session's underlying channel is closed before we call waitForCondition().
         * This means that EOF and STDOUT_DATA (or STDERR_DATA, or both) may
         * be set together.
         */

        int conditions = sess.waitForCondition(ChannelCondition.STDOUT_DATA | 
            ChannelCondition.STDERR_DATA | ChannelCondition.EOF, timeout);

        /* Wait no longer than 2 seconds (= 2000 milliseconds) */
        if ((conditions & ChannelCondition.TIMEOUT) != 0)
        {
          /* A timeout occured. */
          throw new IOException("Timeout while waiting for data from peer");
        }

        /* Here we do not need to check separately for CLOSED, since CLOSED implies EOF */
        if ((conditions & ChannelCondition.EOF) != 0)
        {
          /* The remote side won't send us further data... */
          if ((conditions & (ChannelCondition.STDOUT_DATA | ChannelCondition.STDERR_DATA)) == 0)
          {
            /* ... and we have consumed all data in the local arrival window. */
            break;
          }
        }

        /* OK, either STDOUT_DATA or STDERR_DATA (or both) is set. */

        // You can be paranoid and check that the library is not going nuts:
        // if ((conditions & (ChannelCondition.STDOUT_DATA | ChannelCondition.STDERR_DATA)) == 0)
        //  throw new IllegalStateException("Unexpected condition result (" + conditions + ")");
      }

      /* If you below replace "while" with "if", then the way the output appears on the local
       * stdout and stder streams is more "balanced". Addtionally reducing the buffer size
       * will also improve the interleaving, but performance will slightly suffer.
       * OKOK, that all matters only if you get HUGE amounts of stdout and stderr data =)
       */

      while (stdout.ready())
      {
        int len = stdout.read(buffer);
        builderStdOut.append(buffer, 0, len);
      }

      while (stderr.ready())
      {
        int len = stderr.read(buffer);
        builderStdErr.append(buffer, 0, len);
      }

      if (builderStdOut.toString().contains(prompt))
        break;
    }
  }
  
  public static void main(String[] args) throws IOException {
    WrapSshConnector2 telnet = new WrapSshConnector2();
    telnet.connect("192.168.3.218", 22, "root", "");
    System.out.println(telnet.execute("ls -la"));
    System.out.println(telnet.execute("uname"));
    System.out.println(telnet.execute("true"));
    System.out.println(telnet.execute("false"));
    System.out.println(telnet.execute("ps"));
    telnet.disconnect();
  }

  public class AsynchRunner extends Thread implements JistAPI.DoNotRewrite {
    private int requestId;
    private String[] script;
    private String output = "";
    private IOException e;
    public AsynchRunner(int requestId, String[] script) {
      this.requestId = requestId;
      this.script = script;
    }
    public void run() {
      try {
        for (int i = 0; i < script.length; i++) {
          output += execute(script[i]);
        }
      } catch (IOException e) {
        this.e = e;
      }
    }
  }
  
  public int beginExecute(String[] script) {
    AsynchRunner thread = new AsynchRunner(nextRequestId++, script);
    mapRequests.put(thread.requestId, thread);
    thread.start();
    return thread.requestId;
  }

  public String finishExecute(int requestId) throws IOException {
    AsynchRunner runner = mapRequests.remove(requestId);
    try {
      runner.join();
    } catch (InterruptedException e) {
      // TODO Auto-generated catch block
      e.printStackTrace();
    }
    if (null != runner.e)
      throw runner.e;
    return runner.output;
  }

}
