# include "pcontrol.h"

/* connection can take a while, allow up to 2 sec */
# define CONNECT_TIMEOUT 5000

/* connect to host, start the shell: ssh hostname pclient -> command hostname shell
   stdio is an array of file descriptors (stdio[3])
*/

int rconnect (char *command, char *hostname, char *shell, int *stdio) {

  int i = 0, stdin_fd[2], stdout_fd[2], stderr_fd[2], status;
  int result, waitstatus;
  pid_t pid;
  char *p;
  char **argv;
  IOBuffer buffer;
  struct timespec request, remain;

  ASSERT (command != NULL, "command is NULL");
  ASSERT (hostname != NULL, "hostname is NULL");
  ASSERT (shell != NULL, "shell is NULL");
  ASSERT (stdio != NULL, "stdio is NULL");

  bzero (stdin_fd,  2*sizeof(int));
  bzero (stdout_fd, 2*sizeof(int));
  bzero (stderr_fd, 2*sizeof(int));

  if (pipe (stdin_fd)  < 0) goto pipe_error;
  if (pipe (stdout_fd) < 0) goto pipe_error;
  if (pipe (stderr_fd) < 0) goto pipe_error;

  ALLOCATE (argv, char *, 4);
  argv[0] = command;
  argv[1] = hostname;
  argv[2] = shell;
  argv[3] = 0;

  pid = fork ();
  if (!pid) { /* must be child process */
    if (VerboseMode()) gprint (GP_ERR, "starting remote connection to %s...", hostname);

    /* close the other ends of the pipes */
    close (stdin_fd[1]);
    close (stdout_fd[0]);
    close (stderr_fd[0]);

    /* tie our ends of the pipes to stdin, stdout, stderr */
    dup2 (stdin_fd[0],  STDIN_FILENO);
    dup2 (stdout_fd[1], STDOUT_FILENO);
    dup2 (stderr_fd[1], STDERR_FILENO);

    /* set all three unblocking */
    setvbuf (stdin,  (char *) NULL, _IONBF, BUFSIZ);
    setvbuf (stdout, (char *) NULL, _IONBF, BUFSIZ);
    setvbuf (stderr, (char *) NULL, _IONBF, BUFSIZ);

    status = execvp (argv[0], argv);
    gprint (GP_ERR, "error starting remote shell process\n");
    pcontrol_exit (60);
  }
  free (argv);

  /* close the other ends of the pipes */
  close (stdin_fd[0]);  stdin_fd[0]  = 0;
  close (stdout_fd[1]); stdout_fd[1] = 0;
  close (stderr_fd[1]); stderr_fd[1] = 0;

  /* make the pipes non-blocking */
  fcntl (stdin_fd[1],  F_SETFL, O_NONBLOCK);
  fcntl (stdout_fd[0], F_SETFL, O_NONBLOCK);
  fcntl (stderr_fd[0], F_SETFL, O_NONBLOCK);

  /* perform handshake with pclient to verify alive & running */
  /** this handshake is similar to PclientCommand, but has important differences **/
  InitIOBuffer (&buffer, 0x100);

  /* send handshake command */
  status = write_fmt (stdin_fd[1], "echo CONNECTED\n");
  if ((status == -1) && (errno == EPIPE)) goto connect_error;

  /* try to get evidence connection is alive - wait upto a few seconds */
  p = NULL;
  status = -1;
  for (i = 0; (i < CONNECT_TIMEOUT) && (status != 0) && (p == NULL); i++) {
    status = ReadtoIOBuffer (&buffer, stdout_fd[0]);
    p = memstr (buffer.buffer, "CONNECTED", buffer.Nbuffer);
    usleep (10000); // wait for client to be connected
  }
  if (status == 0) goto connect_error;
  if (status == -1) goto connect_error;
  if (VerboseMode()) gprint (GP_ERR, "%d cycles to connect\n", i);
  FreeIOBuffer (&buffer);

  if (VerboseMode()) gprint (GP_ERR, "Connected\n");

  stdio[0] = stdin_fd[1];
  stdio[1] = stdout_fd[0];
  stdio[2] = stderr_fd[0];

  return (pid);

pipe_error:
  perror ("pipe error:");
  goto close_pipes;

connect_error:
  if (VerboseMode()) gprint (GP_ERR, "error while connecting, status: %d, ncycles: %d\n", status, i);

  /* avoid blocking on waitpid, test every 100 usec, up to 50 msec */
  request.tv_sec = 0;
  request.tv_nsec = 100000;

  /* harvest the child process: kill & wait (< 100 ms) for exit */
  kill (pid, SIGKILL);
  result = waitpid (pid, &waitstatus, WNOHANG);
  for (i = 0; (i < 50) && (result == 0); i++) {
    nanosleep (&request, &remain);
    result = waitpid (pid, &waitstatus, WNOHANG);
  }

  if ((result == -1) && (errno != ECHILD)) {
    gprint (GP_ERR, "unexpected error from waitpid (%d): programming error\n", errno);
    pcontrol_exit (61);
  }
  if (result == 0) {
    if (VerboseMode()) gprint (GP_ERR, "child did not exit (rconnect)??");
  }
  if (result > 0) {
    if (result != pid) {
      gprint (GP_ERR, "waitpid error: mis-matched PID (%d vs %d).  programming error\n", result, pid);
      pcontrol_exit (62);
    }
    if (WIFSTOPPED(waitstatus)) {
      gprint (GP_ERR, "waitpid returns 'stopped': programming error\n");
      pcontrol_exit (63);
    }
  }

close_pipes:
  if (stdin_fd[0]  != 0) close (stdin_fd[0]);
  if (stdin_fd[1]  != 0) close (stdin_fd[1]);
  if (stdout_fd[0] != 0) close (stdout_fd[0]);
  if (stdout_fd[1] != 0) close (stdout_fd[1]);
  if (stderr_fd[0] != 0) close (stderr_fd[0]);
  if (stderr_fd[1] != 0) close (stderr_fd[1]);
  return (FALSE);
}
