/*

Copyright (C) 2019, 2020 Olaf Till <i7tiol@t-online.de>

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 3 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, see <http://www.gnu.org/licenses/>.

*/

#include "parallel-gnutls.h"

#if defined _WIN32 && ! defined __CYGWIN__

#define WIN32_LEAN_AND_MEAN
#include <windows.h>

#else

#include <sys/wait.h>

extern char **environ;

#endif

// The definition of this function is copied from Octave and is in
// 'make-command-string.cc'.
char *make_command_string (const char *cmd, char *const *args);

namespace octave_parallel
{

  instance::instance (int fd, std::string& octb)
    : store_errno (0), pps (),
#if defined _WIN32 && ! defined __CYGWIN__
      process_handle (NULL),
#else
      pid (0),
#endif
      inbuf (NULL), outbuf (NULL),
      is (NULL), os (NULL), cdir (".")
  {
    if (! pps.good ())
      {
        store_errno = errno;
        return;
      }

    outbuf = new octave_parallel_pipe_streambuf (pps.get_out_write (), false);

    os = new std::ostream (outbuf);

    inbuf = new octave_parallel_pipe_streambuf (pps.get_in_read (), true);

    is = new std::istream (inbuf);

    char *args [4];
    copy_to_non_const arg_binary (octb.c_str ());
    args[0] = arg_binary.get_str ();
    copy_to_non_const arg_eval_option ("--eval");
    args[1] = arg_eval_option.get_str ();
    std::stringstream eval_arg;
    eval_arg << "pkg load parallel; __rfeval_slave__("
             << (uint64_t)pps.get_out_read_handle () << ", "
             << (uint64_t)pps.get_in_write_handle ()
             << ")";
    std::string eval_arg_str (eval_arg.str ());
    copy_to_non_const arg_eval_command (eval_arg_str.c_str ());
    args[2] = arg_eval_command.get_str ();
    args[3] = NULL;

#if defined _WIN32 && ! defined __CYGWIN__

    char *command = make_command_string (args[0], args);

    STARTUPINFO si;
    ZeroMemory (&si, sizeof (si));
    si.cb = sizeof (si);

    PROCESS_INFORMATION pi;
    ZeroMemory (&pi, sizeof (pi));

    if (CreateProcessA (NULL, command, NULL, NULL, TRUE, 0, NULL, NULL,
                        &si, &pi))
      {
        process_handle = pi.hProcess;
        CloseHandle (pi.hThread);

        pps.close_child_side ();

        if (! pps.good ())
          {
            _p_error ("could not close child pipe ends in parent");
            store_errno = -1;
          }
      }
    else
      {
        _p_error ("process generation failed");
        store_errno = -1;
      }

    free (command);

#else // Posix API

    char **envp = environ;

    pid_t tp;

    if ((tp = ::fork ()) == 0)
      {
        // child process

        pps.close_parent_side ();

        if (! pps.good ())
          {
            _p_error ("could not close parent pipe ends in child");
            _exit (1);
          }

#ifdef HAVE_FEXECVE

        if (fexecve (fd, args, envp))
          {
            gnulib_perror ("fexecve");
            _exit (1);
          }

#else

        if (execve (args[0], args, envp))
          {
            gnulib_perror ("execve");
            _exit (1);
          }

#endif

      }
    else if (tp == -1)
      {
        // parent process, error in process generation
        _p_error ("process generation failed");
        store_errno = errno;
      }
    else
      {
        // parent process after successful process generation

        pid = tp;

        pps.close_child_side ();

        if (! pps.good ())
          {
            _p_error ("could not close child pipe ends in parent");
            store_errno = -1;
          }
      }

#endif
  }

  instance::~instance (void)
  {
    // send termination signal
    octave_value ts = octave_value (octave_int32 (1));
    send_octave_value (ts);

    delete is;

    delete os;

    delete inbuf;

    delete outbuf;

#if defined _WIN32 && ! defined __CYGWIN__

    if (WaitForSingleObject (process_handle, 1) == WAIT_TIMEOUT
        && WaitForSingleObject (process_handle, 5) == WAIT_TIMEOUT)
      {
        if (TerminateProcess (process_handle, 0))
          {
            // Is this necessary on Windows?
            WaitForSingleObject (process_handle, INFINITE);
          }
        else
          _p_error ("could not terminate child process");
      }

    CloseHandle (process_handle);

#else // Posix API

    usleep (1000);
    if (! waitpid (pid, NULL, WNOHANG))
      {
        usleep (5000);
        if (! waitpid (pid, NULL, WNOHANG))
          {
            kill (pid, SIGKILL);
            waitpid (pid, NULL, 0);
          }
      }

#endif
  }

  int instance::initialize_new_function_handle
  (octave_value& fcn, std::string& dir, octave_value& eh, octave_value& nargout)
  {
    if (send_octave_value (fcn))
      return -1;

    octave_value ov_dir = dir;
    if (send_octave_value (ov_dir))
      return -1;

    octave_value result;

    if (recv_octave_value (result))
      {
        return -1;
      }

    if (! result.OV_ISINTEGER ())
      {
        store_errno = -1;
        return -1;
      }

    if (! result.int_value ())
      {
        store_errno = -1;
        error ("could not change directory in slave");
        return -1;
      }

    cdir = dir;

    octave_value nout = octave_value (octave_int32 (nargout.int_value ()));

    if (eh.is_defined ())
      {
        if (send_octave_value (eh))
          return -1;

        if (send_octave_value (nout))
          return -1;
      }
    else
      {
        if (send_octave_value (nout))
          return -1;
      }

    return 0;
  }

  control::control (const std::string& path)
    : nproc (0), nproc_used (0), nproc_max (0), octave_fd (-1),
      unrecoverable_error (0), instances (), states (), posix_fds (),
      gnulib_pollfds (NULL), poll_res (), fcn (), eh (), dir (),
      nargout ()
    {
      if (set_executable (path))
        {
          _p_error ("couldn't set executable");

          unrecoverable_error = 1;
        }
      else
        {
          nproc_max = num_processors (NPROC_CURRENT);

          gnulib_pollfds = gnulib_alloc_pollfds (nproc_max);
        }
    }        

  int control::set_nproc (octave_idx_type np)
    {
      np = np > nproc_max ? nproc_max : np;

      if (np < 0)
        {
          _p_error ("internal error, negative process count requested");

          return -1;
        }
      else if (np > nproc)
        {
          for (octave_idx_type id = nproc; id < np; id++)
            {
              instance *tp_instc = new instance (octave_fd, octave_binary);

              if (! tp_instc->good ())
                {
                  delete tp_instc;

                  _p_error ("couldn't create process");

                  return -1;
                }

              if (fcn.is_defined ()
                  && tp_instc->initialize_new_function_handle
                  (fcn, dir, eh, nargout))
                {
                  delete tp_instc;

                  _p_error ("couldn't initialize new process");

                  return -1;
                }

              instances.push_back (tp_instc);

              states.push_back (0);

              posix_fds.push_back (tp_instc->get_in_read_posix_fd ());

              poll_res.push_back (0);

            }
        }
      else if (np < nproc)
        {
          for (octave_idx_type id = nproc - 1; id >= np; id--)
            {
              delete instances[id];

              instances.pop_back ();

              states.pop_back ();

              posix_fds.pop_back ();

              poll_res.pop_back ();
            }
        }
      else
        {
          // np == nproc, do nothing
        }

      nproc = np;

      if (nproc < nproc_used)
        {
          nproc_used = nproc;
        }

      return 0;
    }

  int control::set_nproc_used (octave_idx_type npu)
    {
      npu = npu > nproc_max ? nproc_max : npu;

      if (npu < 0)
        {
          _p_error ("internal error, negative count of used processes requested");

          return -1;
        }
      else
        {
          if (npu > nproc && set_nproc (npu))
            return -1;
          else
            nproc_used = npu;
        }

      return 0;
    }

  int control::set_executable (const std::string& path)
  {
    if (unrecoverable_error)
      return -1;

    std::string tp_octave_binary;

    if (path.empty ())
      {
        char *cbindir = ::getenv ("OCTAVE_BINDIR");

        std::string octave_bindir = cbindir ? cbindir : OCTAVE__CONFIG__BIN_DIR;

        tp_octave_binary = octave_bindir + OCTAVE__SYS__FILE_OPS::dir_sep_str ()
          + "octave-cli";

#if defined _WIN32 && ! defined __CYGWIN__
        tp_octave_binary = tp_octave_binary + ".exe";
#endif
     }
    else
      {
        tp_octave_binary = path;
      }

#if defined _WIN32 && ! defined __CYGWIN__
    octave_binary = tp_octave_binary;
#else
    // open the binary to avoid racing conditions
    int tp_fd = ::open (tp_octave_binary.c_str (), O_RDONLY | O_CLOEXEC);

    if (tp_fd == -1)
      {
        gnulib_perror ("open");

        return -1;
      }
    else if (gnulib_fstat_reg_exec_nonsuid (tp_fd))
      {
        _p_error ("couldn't verify that Octave binary is an executable regular file that hasn't the set-user-ID bit set");
        ::close (tp_fd);

        return -1;
      }
    else
      {
        octave_binary = tp_octave_binary;
        if (octave_fd != -1)
          ::close (octave_fd);
        octave_fd = tp_fd;
      }
#endif

    return 0;
  }

  int control::poll_input (int timeout)
  {

    int size = posix_fds.size ();

    for (int id = 0; id < size; id++)
      {
        poll_res[id] = posix_fds[id];
      }

    return gnulib_poll_in_wrapper (poll_res.data (), gnulib_pollfds,
                                   size, timeout);
  }

  int control::initialize_job
  (const octave_value& a_fcn, const std::string& a_dir,
   const octave_value& a_eh, const octave_value& a_nargout)
  {
    if (unrecoverable_error)
      return -1;

    fcn = a_fcn;
    dir = a_dir;
    eh = a_eh;
    nargout = a_nargout;

    for (int id = 0; id < nproc; id++)
      {
        if (states[id] > 0
            || instances[id]->initialize_new_function_handle
            (fcn, dir, eh, nargout))
          {
            unrecoverable_error = 1;

            return -1;
          }
      }

    return 0;
  }

  octave_value control::get_next_result (octave_value& res)
  {
    octave_value retval;

    if (unrecoverable_error)
      return retval;

    int poll_ret = poll_input (-1);

    int id_ready = -1;

    if (poll_ret < 1)
      {
        unrecoverable_error = 1;
      }
    else
      {
        int size = posix_fds.size ();

        for (int id = 0; id < size; id++)
          {
            if (poll_res[id])
              {
                id_ready = id;

                break;
              }
          }

        if (id_ready == -1)
          {
            unrecoverable_error = 1;
          }
      }

    if (unrecoverable_error)
      {
        _p_error ("poll_error");

        return retval;
      }

    if (instances[id_ready]->recv_octave_value (res))
      {
        unrecoverable_error = 1;

        _p_error ("could not receive result");

        return retval;
      }

    if (res.isinteger ()
        && res.int_value () == -2)
      {
        unrecoverable_error = 1;

        _p_error ("error in error handler");

        return retval;
      }

    // 1-based index
    retval = octave_value (states[id_ready]);

    states[id_ready] = 0;

    return retval;
  }

  int control::send_next_args (octave_value& idx,
                               octave_value& args)
  {
    if (unrecoverable_error)
      return -1;

    int ret = 0;

    for (int id = 0; id < nproc_used; id++)
      {
        if (states[id] == 0)
          {
            if (instances[id]->send_octave_value (args))
              {
                unrecoverable_error = 1;

                _p_error ("could not send arguments");

                ret = -1;
              }
            else if (eh.is_defined ()
                     && instances[id]->send_octave_value (idx))
              {
                unrecoverable_error = 1;

                _p_error ("could not send argument index");

                ret = -1;
              }
            else
              {
                // 1-based index
                states[id] = idx.idx_type_value ();

                ret = 1;
              }

            break;
          }
      }

    return ret;
  }

} // namespace octave_parallel
