// 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
//==============================================================================

#include <iostream> // Temporary.
#include <stdexcept>

#include "child.hpp"

namespace CodeDweller {

  Child::Child(std::vector<std::string> args, size_t bufSize) :
    readStreambuf(bufSize),
    writeStreambuf(bufSize),
    reader(&readStreambuf),
    writer(&writeStreambuf) {

    init();

    if (args.size() == 0) {
      //
    } else if (args.size() == 1) {
      cmdline = args[0];
      return;
    }

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

  Child::Child(std::string childpath, size_t bufSize) :
    readStreambuf(bufSize),
    writeStreambuf(bufSize),
    reader(&readStreambuf),
    writer(&writeStreambuf),
    cmdline(childpath) {
    init();
  }

  Child::~Child() {
    // Close handles.
  }

  void
  Child::init() {
    childStarted = false;
    exitCodeObtainedFlag = false;
    exitCode = 0;
  }

  void
  Child::run() {

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

    // 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;
 
   // Create the child process. 
   bool status;
 
   status = CreateProcess(NULL, 
			  (char *) cmdline.c_str(), // command line 
			  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.
   readStreambuf.setInputHandle(childStdOutAtParent);
   writeStreambuf.setOutputHandle(childStdInAtParent);

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

   childStarted = true;

   // 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());
    }

  }

  void
  Child::terminate() {

    if (isDone()) {

      return;

    }

    if (!TerminateProcess(childProcess, terminateExitCode)) {
      throw std::runtime_error("Error terminating the child process:  " +
			       getErrorText());
    }

  }

  bool
  Child::isDone() {

    if (exitCodeObtainedFlag) {

      return true;

    }

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

    int result;

    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;

    return true;

  }

  int32_t
  Child::result() {

    if (exitCodeObtainedFlag) {

      return exitCode;

    }

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

    int result;

    if (!GetExitCodeProcess(childProcess, (LPDWORD) &result)) {
      throw std::runtime_error("Error getting child process exit code:  " +
			       getErrorText());
    }

    // Error if the process has not exited.
    if (STILL_ACTIVE == result) {
      throw std::logic_error("Child process was active when "
			     "result() was called");
    }

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

    return result;
  }

  std::string
  Child::getErrorText() {

    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;
  }

  Child::ReadStreambuf::ReadStreambuf(std::size_t bufferSize) :
    inputHandle(0),
    buffer(bufferSize + 1) {

    char *end = &(buffer.front()) + buffer.size();

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

  }

  void
  Child::ReadStreambuf::setInputHandle(HANDLE inHandle) {

    inputHandle = inHandle;

  }

  std::streambuf::int_type
  Child::ReadStreambuf::underflow() {

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

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

    }

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

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

      // Not the first fill.  Copy putback characters.

    }

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

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

    // 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());

  }

  Child::WriteStreambuf::WriteStreambuf(std::size_t bufferSize) :
    outputHandle(0),
    buffer(bufferSize + 1) {

    char *base = &(buffer.front());

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

  }

  void
  Child::WriteStreambuf::setOutputHandle(HANDLE outHandle) {

    outputHandle = outHandle;

  }

  void
  Child::WriteStreambuf::flushBuffer() {

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

    if (!WriteFile(outputHandle,
		   pbase(),
		   nBytes,
		   &nBytesWritten,
		   NULL)) {
      throw std::runtime_error("Error writing to child process:  " +
			       getErrorText());
    }
    if (nBytes != nBytesWritten) {
      throw std::runtime_error("Not all data was written to to child process:  " +
			       getErrorText());
    }
    pbump(-nBytes);

    return;

  }

  std::streambuf::int_type
  Child::WriteStreambuf::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 traits_type::not_eof('a');

  }

  int
  Child::WriteStreambuf::sync() {

    flushBuffer();
    return 1; // Success.

  }

}