// \file child.cpp
//
// Copyright (C) 2014 MicroNeil Research Corporation.
//
// This program is part of the MicroNeil Research Open Library Project. For
// more information go to http://www.microneil.com/OpenLibrary/index.html
//
// This program is free software; you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by the
// Free Software Foundation; either version 2 of the License, or (at your
// option) any later version.
//
// This program is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
// more details.
//
// You should have received a copy of the GNU General Public License along with
// this program; if not, write to the Free Software Foundation, Inc., 59 Temple
// Place, Suite 330, Boston, MA 02111-1307 USA
//==============================================================================

#ifndef _WIN32
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <signal.h>
#include <sys/select.h>

#include <cstring>
#include <cerrno>
#endif

#include <stdexcept>

#include "CodeDweller/child.hpp"

namespace CodeDweller {

  ChildStream::ChildStream(std::vector<std::string> const &args,
                           size_t bufSize) :
    childStreambuf(bufSize),
    std::iostream(&childStreambuf),
    cmdArgs(args) {

    init();
    run();

  }

  ChildStream::ChildStream(std::string const &childpath, size_t bufSize) :
    childStreambuf(bufSize),
    std::iostream(&childStreambuf) {

    cmdArgs.push_back(childpath);
    init();
    run();
  }

  ChildStream::ChildStream(size_t bufSize) :
    childStreambuf(bufSize),
    std::iostream(&childStreambuf) {

    init();
  }

  ChildStream::~ChildStream() {
    if (isRunning()) {
      close();
    }
  }

  void ChildStream::init() {

    childStarted = false;
    childExited = false;
    exitCodeObtainedFlag = false;
    exitCode = 0;
  }

  void ChildStream::open(std::vector<std::string> const &args) {
    cmdArgs = args;
    init();
    run();
  }

  void ChildStream::open(std::string const &childpath) {
    cmdArgs.clear();
    cmdArgs.push_back(childpath);
    init();
    run();
  }

  size_t ChildStream::numBytesAvailable() const {

    return childStreambuf.numBytesAvailable();

  }

  bool ChildStream::isRunning() const {
    return childStarted && !childExited;
  }

  void ChildStream::run() {

    if (childStarted) {
      throw std::logic_error("Child process was active when "
			     "run() was called");
    }

    if (cmdArgs.empty()) {
      throw std::invalid_argument("A child executable must be specified.");
    }

#ifdef _WIN32
    // Set the bInheritHandle flag so pipe handles are inherited. 
    SECURITY_ATTRIBUTES securityAttributes; 

    securityAttributes.nLength = sizeof(SECURITY_ATTRIBUTES); 
    securityAttributes.bInheritHandle = true; 
    securityAttributes.lpSecurityDescriptor = NULL; 

    // Create a pipe for the child process's STDOUT.
    HANDLE childStdOutAtChild;
    HANDLE childStdOutAtParent;
    HANDLE childStdInAtChild;
    HANDLE childStdInAtParent;
    int bufferSize = 0;

    if (!CreatePipe(&childStdOutAtParent,
		    &childStdOutAtChild,
		    &securityAttributes,
		    bufferSize)) {
      throw std::runtime_error("Error from CreatePipe for stdout:  " +
			       getErrorText());
    }

    // Ensure the read handle to the pipe for STDOUT is not inherited.
    int inheritFlag = 0;
    if (!SetHandleInformation(childStdOutAtParent,
			      HANDLE_FLAG_INHERIT,
			      inheritFlag) ) {
      throw std::runtime_error("Error from GetHandleInformation for stdout:  " +
			       getErrorText());
    }

    // Create a pipe for the child process's STDIN. 
    if (! CreatePipe(&childStdInAtChild,
		     &childStdInAtParent,
		     &securityAttributes,
		     bufferSize)) {
      throw std::runtime_error("Error from CreatePipe for stdin:  " +
			       getErrorText());
    }

    // Ensure the write handle to the pipe for STDIN is not inherited. 
    if (!SetHandleInformation(childStdInAtParent,
			      HANDLE_FLAG_INHERIT,
			      inheritFlag)) {
      throw std::runtime_error("Error from GetHandleInformation for stdin:  " +
			       getErrorText());
    }
 
    // Set up members of the PROCESS_INFORMATION structure. 
    PROCESS_INFORMATION processInfo; 

    std::fill((char *) &processInfo,
	      ((char *) &processInfo) + sizeof(PROCESS_INFORMATION),
	      0);
 
    // Set up members of the STARTUPINFO structure.  This structure
    // specifies the STDIN and STDOUT handles for redirection.
    STARTUPINFO startInfo;

    std::fill((char *) &startInfo,
	      ((char *) &startInfo) + sizeof(STARTUPINFO),
	      0);
    startInfo.cb = sizeof(STARTUPINFO); 
    startInfo.hStdError = childStdOutAtChild;
    startInfo.hStdOutput = childStdOutAtChild;
    startInfo.hStdInput = childStdInAtChild;
    startInfo.dwFlags |= STARTF_USESTDHANDLES;
 
    // Assemble the command line.
    std::string cmdline;

    if (cmdArgs.size() == 1) {
      cmdline = cmdArgs[0];
    } else {

      // Append all but last command-line arguments.
      for (size_t i = 0; i < cmdArgs.size() - 1; i++) {
	cmdline += cmdArgs[i] + " ";
      }
      cmdline += cmdArgs.back(); // Append last command-line argument.
    }

    // Create the child process. 
    bool status;
 
    status = CreateProcess(NULL, 
			   (char *) cmdline.c_str(),
			   NULL,          // process security attributes 
			   NULL,          // primary thread security attributes 
			   true,          // handles are inherited 
			   0,             // creation flags 
			   NULL,          // use parent's environment 
			   NULL,          // use parent's current directory 
			   &startInfo,    // STARTUPINFO pointer 
			   &processInfo); // receives PROCESS_INFORMATION 

    // If an error occurs, exit the application. 
    if (!status ) {
      throw std::runtime_error("Error from CreateProcess with "
			       "command line \"" + cmdline + "\":  " +
			       getErrorText());
    }

    // Provide the stream buffers with the handles for communicating
    // with the child process.
    childStreambuf.setInputHandle(childStdOutAtParent);
    childStreambuf.setOutputHandle(childStdInAtParent);

    // Save the handles to the child process and its primary thread.
    childProcess = processInfo.hProcess;
    childThread = processInfo.hThread;

    // Close the child's end of the pipes.
    if (!CloseHandle(childStdOutAtChild)) {
      throw std::runtime_error("Error closing the child process "
			       "stdout handle:  " + getErrorText());
    }
     
    if (!CloseHandle(childStdInAtChild)) {
      throw std::runtime_error("Error closing the child process "
			       "stdin handle:  " + getErrorText());
    }

#else

    // Create the pipes for the stdin and stdout.
    int childStdInPipe[2];
    int childStdOutPipe[2];

    if (pipe(childStdInPipe) != 0) {
      throw std::runtime_error("Error creating pipe for stdin:  " +
			       getErrorText());
    }
    if (pipe(childStdOutPipe) != 0) {
      ::close(childStdInPipe[0]);
      ::close(childStdInPipe[1]);
      throw std::runtime_error("Error creating pipe for stdout:  " +
			       getErrorText());
    }

    // Create the child process. 
    childPid = fork();
    if (-1 == childPid) {
      for (int i = 0; i < 2; i++) {
	::close(childStdInPipe[i]);
	::close(childStdOutPipe[i]);
      }
      throw std::runtime_error("Error creating child process:  " +
			       getErrorText());
    }
 
    if (0 == childPid) {

      // The child executes this.  Redirect stdin.
      if (dup2(childStdInPipe[0], STDIN_FILENO) == -1) {
	std::string errMsg;

	// Send message to parent.
	errMsg = "Error redirecting stdin in the child:  " + getErrorText();
	::write(childStdOutPipe[1], errMsg.data(), errMsg.size());
	exit(-1);
      }

      // Redirect stdout.
      if (dup2(childStdOutPipe[1], STDOUT_FILENO) == -1) {
	std::string errMsg;

	// Send message to parent.
	errMsg = "Error redirecting stdout in the child:  " + getErrorText();
	::write(childStdOutPipe[1], errMsg.data(), errMsg.size());
	exit(-1);
      }

      // Close pipes.
      if ( (::close(childStdInPipe[0]) != 0) ||
	   (::close(childStdInPipe[1]) != 0) ||
	   (::close(childStdOutPipe[0]) != 0) ||
	   (::close(childStdOutPipe[1]) != 0) ) {
	std::string errMsg;

	// Send message to parent.
	errMsg = "Error closing the pipes in the child:  " + getErrorText();
	::write(STDOUT_FILENO, errMsg.data(), errMsg.size());
	exit(-1);

      }

      // Prepare the arguments.
      std::vector<const char *> execvArgv;

      for (auto &arg : cmdArgs) {
	execvArgv.push_back(arg.c_str());
      }
      execvArgv.push_back((char *) NULL);

      // Run the child process image.
      (void) execv(execvArgv[0], (char * const *) &(execvArgv[0]));

      // Error from exec.
      std::string errMsg;

      // Send message to parent.
      errMsg = "Error from exec:  " + getErrorText();
      ::write(STDOUT_FILENO, errMsg.data(), errMsg.size());
      exit(-1);

    }

    // Provide the stream buffers with the file descriptors for
    // communicating with the child process.
    childStreambuf.setInputFileDescriptor(childStdOutPipe[0]);
    childStreambuf.setOutputFileDescriptor(childStdInPipe[1]);

    // Close the child's end of the pipes.
    if ( (::close(childStdInPipe[0]) != 0) ||
	 (::close(childStdOutPipe[1]) != 0) ) {
      std::string errMsg;

      throw std::runtime_error("Error closing child's end of pipes in "
			       "the parent:  " + getErrorText());
    }

#endif

    childStarted = true;

  }

  void ChildStream::close() {

    if (isDone()) {

      return;

    }

#ifdef _WIN32
    if (!TerminateProcess(childProcess, terminateExitCode)) {
#else
    if (kill(childPid, SIGTERM) != 0) {
#endif
      throw std::runtime_error("Error terminating the child process:  " +
			       getErrorText());
    }

  }

  bool ChildStream::isDone() {

    if (childExited) {

      return true;

    }

    if (!childStarted) {
      throw std::logic_error("Child process was not started "
			     "when isDone() was called");
    }

    int result;

#ifdef _WIN32
    if (!GetExitCodeProcess(childProcess, (LPDWORD) &result)) {
      throw std::runtime_error("Error checking status of child process:  " +
			       getErrorText());
    }

    if (STILL_ACTIVE == result) {

      return false;

    }

    // Child process has exited.  Save the exit code.
    exitCode = result;
    exitCodeObtainedFlag = true;

#else
    int status = 0;

    result = waitpid(childPid, &status, WNOHANG);
    if (-1 == result) {
      throw std::runtime_error("Error checking status of child process:  " +
			       getErrorText());
    } else if (0 == result) {

      // Child is still running.
      return false;

    }

    if (WIFEXITED(status)) {

      // Child exited normally.
      exitCode = WEXITSTATUS(status);
      exitCodeObtainedFlag = true;

    }

#endif

    childExited = true;

    return true;

  }

  int32_t ChildStream::result() {

    if (exitCodeObtainedFlag) {

      return exitCode;

    }

    // Check whether the process is running, and get the exit code.
    if (!isDone()) {
      throw std::logic_error("Child process was still running "
			     "when result() was called");
    }

    // Child process has exited.
    if (!exitCodeObtainedFlag) {

      // Exit code is not available.
      throw std::runtime_error("Child process has exited but the exit "
			       "code is not available");

    }

    return exitCode;

  }

  std::string ChildStream::getErrorText() {

#ifdef _WIN32
    LPVOID winMsgBuf;
    DWORD lastError = GetLastError(); 

    FormatMessage(
		  FORMAT_MESSAGE_ALLOCATE_BUFFER | 
		  FORMAT_MESSAGE_FROM_SYSTEM |
		  FORMAT_MESSAGE_IGNORE_INSERTS,
		  NULL,
		  lastError,
		  MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
		  (char *) &winMsgBuf,
		  0, NULL );

    std::string errMsg((char *) winMsgBuf);

    LocalFree(winMsgBuf);
    return errMsg;
#else
    return strerror(errno);
#endif
  }

  ChildStream::ChildStreambuf::ChildStreambuf(std::size_t bufSize) :
#ifdef _WIN32
    inputHandle(0),
    outputHandle(0),
#else
    inputFileDescriptor(-1),
    outputFileDescriptor(-1),
#endif
    bufferSize(bufSize),
    readBuffer(bufferSize + 1),
    writeBuffer(bufferSize + 1) {

    // Read buffer initialization.
    char *end = &(readBuffer.front()) + readBuffer.size();

    // Indicate to underflow that underflow has not been called.
    setg(end, end, end);

    // Write buffer initialization.
    char *base = &(writeBuffer.front());

    // Indicate to overflow that overflow has not been called.
    setp(base, base + writeBuffer.size() - 1);

  }

#ifdef _WIN32
  void ChildStream::ChildStreambuf::setInputHandle(HANDLE inHandle) {

    inputHandle = inHandle;

  }
#else
  void ChildStream::ChildStreambuf::setInputFileDescriptor(int inFd) {

    inputFileDescriptor = inFd;

  }
#endif

  size_t ChildStream::ChildStreambuf::numBytesAvailable() const {

    size_t nBytesAvailable =  egptr() - gptr();

#ifdef _WIN32
    DWORD lpTotalBytesAvail;

    if (!PeekNamedPipe(inputHandle,
		       NULL,
		       0,
		       NULL,
		       &lpTotalBytesAvail,
		       NULL)) {
      throw std::runtime_error("Error from PeekNamedPipe:  " +
			       getErrorText());
    }

    if (lpTotalBytesAvail > 0) {
      nBytesAvailable++;
    }
#else
    fd_set readFd;
    int retVal;
    struct timeval timeout = {0, 0};

    FD_ZERO(&readFd);
    FD_SET(inputFileDescriptor, &readFd);

    // Check if input is available.
    retVal = select(inputFileDescriptor + 1, &readFd, NULL, NULL, &timeout);

    if (-1 == retVal) {
      throw std::runtime_error("Error from select(): " + getErrorText());
    } else if (retVal > 0) {
	nBytesAvailable++;
    }
#endif

    return nBytesAvailable;

  }

  std::streambuf::int_type ChildStream::ChildStreambuf::underflow() {

    // Check for empty buffer.
    if (gptr() < egptr()) {

      // Not empty.
      return traits_type::to_int_type(*gptr());

    }

    // Need to fill the buffer.
    char *base = &(readBuffer.front());
    char *start = base;

    // Check whether this is the first fill.
    if (eback() == base) {

      // Not the first fill.  Copy one putback character.
      *(eback()) = *(egptr() - 1);
      start++;
    }

    // start points to the start of the buffer.  Fill buffer.
#ifdef _WIN32
    DWORD nBytesRead;

    if (!ReadFile(inputHandle,
		  start,
		  readBuffer.size() - (start - base),
		  &nBytesRead,
		  NULL)) {
      return traits_type::eof();
    }
#else
    ssize_t nBytesRead;

    nBytesRead = ::read(inputFileDescriptor,
                        start,
                        readBuffer.size() - (start - base));
    if (-1 == nBytesRead) {
      return traits_type::eof();
    }
#endif

    // Check for EOF.
    if (0 == nBytesRead) {
      return traits_type::eof();
    }

    // Update buffer pointers.
    setg(base, start, start + nBytesRead);

    return traits_type::to_int_type(*gptr());

  }

#ifdef _WIN32
  void ChildStream::ChildStreambuf::setOutputHandle(HANDLE outHandle) {

    outputHandle = outHandle;

  }
#else
  void ChildStream::ChildStreambuf::setOutputFileDescriptor(int outFd) {

    outputFileDescriptor = outFd;

  }
#endif

  void ChildStream::ChildStreambuf::flushBuffer() {

    // Write.
    std::ptrdiff_t nBytes = pptr() - pbase();

#ifdef _WIN32
    DWORD nBytesWritten;

    if (!WriteFile(outputHandle,
		   pbase(),
		   nBytes,
		   &nBytesWritten,
		   NULL)) {
      // Clear the output buffer.
      pbump(-nBytes);
      throw std::runtime_error("Error writing to child process:  " +
			       getErrorText());
    }
#else
    ssize_t nBytesWritten;

    nBytesWritten = ::write(outputFileDescriptor, pbase(), nBytes);

#endif

    // Clear the output buffer.
    pbump(-nBytes);

    if (nBytes != nBytesWritten) {
      throw std::runtime_error("Not all data was written to to child "
			       "process:  " + getErrorText());
    }

    return;

  }

  std::streambuf::int_type ChildStream::ChildStreambuf::overflow(int_type ch) {

    // Check whether we're writing EOF.
    if (traits_type::eof() != ch) {

      // Not writing EOF.
      *(pptr()) = ch;
      pbump(1);

      // Write.
      flushBuffer();

      // Success.
      return ch;

    }

    return traits_type::eof();

  }

  int ChildStream::ChildStreambuf::sync() {

    flushBuffer(); // Throws exception on failure.

    // Success.
    return 1;

  }

}