#include <Python.h>


#ifdef WITH_NEXT_FRAMEWORK
/* On Darwin/MacOSX a shared library or framework has no access to
** environ directly, we must obtain it with _NSGetEnviron().
*/
#include <crt_externs.h>
static char **environ;
#elif !defined(_MSC_VER) && ( !defined(__WATCOMC__) || defined(__QNX__) )
extern char **environ;
#endif /* !_MSC_VER */

/* needed for fork(), setsid(), and execve() */
#ifdef HAVE_UNISTD_H
#include <unistd.h>
#endif

#ifdef HAVE_FTIME
#include <sys/timeb.h>
#endif

#ifdef HAVE_SYS_WAIT_H
#include <sys/wait.h>  /* For waitpid() and WNOHANG */
#endif

#ifdef MS_WINDOWS
#include <windows.h>
#include <process.h>
#endif

/* abstract differences between Python versions */
#if PY_MAJOR_VERSION == 1 && PY_MINOR_VERSION < 6
#define PyObject_New PyObject_NEW
#define PyObject_Del PyMem_Del
#endif

typedef struct {
  PyObject_HEAD
  char *  ps_path;
  PyObject *ps_argv;
  PyObject *ps_env;
  int     ps_pid;
#ifdef MS_WINDOWS
  HANDLE  ps_handle;
#endif
} PyProcessObject;

#define PyProcess_GET_PATH(ob) (((PyProcessObject *)(ob))->ps_path)
#define PyProcess_GET_ARGV(ob) (((PyProcessObject *)(ob))->ps_argv)
#define PyProcess_GET_ENV(ob) (((PyProcessObject *)(ob))->ps_env)


/* Global variable holding the exception type for errors detected
   by this module (but not argument type or memory errors, etc.). */

static PyObject *PyProcess_Error;

/* forwards */

staticforward PyTypeObject PyProcess_Type;

/* PyProcessObject functions */

static char **create_environ(PyObject *mapping)
{
  char **envlist;
  PyObject *items;
  int envc, i, j;

  envc = PyMapping_Size(mapping);
  envlist = PyMem_NEW(char *, envc + 1);
  if (envlist == NULL)
  {
    PyErr_NoMemory();
    return NULL;
  }

  items = PyMapping_Items(mapping);
  if (!items)
  {
    PyMem_DEL(envlist);
    return NULL;
  }

  for (i = 0; i < envc; i++)
  {
    char *key, *val, *envset;
    int keylen, vallen;
    PyObject *tuple = PyList_GetItem(items, i);
    if (!tuple) {
      Py_DECREF(items);
      while (--i >= 0)
        PyMem_DEL(envlist[i]);
      PyMem_DEL(envlist);
      return NULL;
    }

    if (!PyArg_ParseTuple(tuple, "s#s#", &key, &keylen, &val, &vallen)) {
      Py_DECREF(tuple);
      Py_DECREF(items);
      while (--i >= 0)
        PyMem_DEL(envlist[i]);
      PyMem_DEL(envlist);
      return NULL;
    }

    envset = PyMem_NEW(char, keylen + vallen + 2);
    if (envset == NULL) {
      PyErr_NoMemory();
      while (--i >= 0)
        PyMem_DEL(envlist[i]);
      PyMem_DEL(envlist);
      Py_DECREF(tuple);
      Py_DECREF(items);
      return NULL;
    }
    sprintf(envset, "%s=%s", key, val);

    /* place this pair in alphabetical order in the envlist */
    /* required for Win32 CreateProcess() but doesn't hurt execve() */
    j = i;
    while (j > 0 && strcmp(envset, envlist[j - 1]) < 0) {
      envlist[j] = envlist[j - 1];
      j--;
    }
    envlist[j] = envset;
  }
  envlist[i] = NULL;

  return envlist;
}

static void free_block(char **block)
{
  int i;
  for (i = 0; block[i]; i++)
    PyMem_DEL(block[i]);
  PyMem_DEL(block);
}

static PyObject *create_process(PyProcessObject *ps)
{
#ifdef MS_WINDOWS
  STARTUPINFO si;
  PROCESS_INFORMATION pi;
  char commandline[512];
  char *cp, *envblock;
  int i;

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

  /* construct the command line */
  cp = commandline;
  cp += sprintf(commandline, "\"%s\"", ps->ps_path);
  for (i = 1; i < PyTuple_GET_SIZE(ps->ps_argv); i++)
  {
    PyObject *arg = PyTuple_GetItem(ps->ps_argv, i);
    if (arg == NULL) return NULL;
    cp += sprintf(cp, " \"%s\"", PyString_AsString(arg));
  }

  if (ps->ps_env)
  {
    int size = 1;
    char *e;
    char **envlist = create_environ(ps->ps_env);
    if (envlist == NULL) return NULL;
    for (i = 0; envlist[i]; i++) size += strlen(envlist[i]) + 1;
    envblock = e = (char *)malloc(size);
    if (envblock == NULL) {
      free_block(envlist);
      return NULL;
    }
    for (i = 0; envlist[i]; i++) {
      strcpy(e, envlist[i]);
      e += strlen(envlist[i]) + 1;
    }
    *e = '\0';
    free_block(envlist);
  }
  else
    envblock = NULL;

  if (!CreateProcess(NULL, commandline, NULL, NULL,
                     TRUE,         /* inherit handles */
                     0, /* create flags */
                     envblock, NULL,
                     &si, &pi))
  {
    PyErr_SetString(PyProcess_Error, "unable to create new process");
    /*
     * We must close the handles to the new process and its main thread
     * to prevent handle and memory leaks.
     */
    if (envblock) 
      free(envblock);
    CloseHandle(pi.hProcess);
    CloseHandle(pi.hThread);
    return NULL;
  }
  if (envblock)
    free(envblock);
  CloseHandle(pi.hThread);
  ps->ps_handle = pi.hProcess;
  ps->ps_pid = pi.dwProcessId;

#else /* not Windows */
  char **argv, **env;
  int i, argc, pid;

  /* construct the argument block */
  argc = PyTuple_Size(ps->ps_argv);
  argv = PyMem_NEW(char *, argc + 1);
  if (argv == NULL) {
    PyErr_NoMemory();
    return NULL;
  }
  for (i = 0; i < argc; i++) {
    PyObject *arg = PyTuple_GetItem(ps->ps_argv, i);
    if (arg == NULL)
    {
      PyMem_DEL(argv);
      return NULL;
    }
    argv[i] = PyString_AsString(arg);
    if (argv[i] == NULL)
    {
      PyMem_DEL(argv);
      return NULL;
    }
  }
  argv[argc] = NULL;

  /* construct the environment block */
#ifdef WITH_NEXT_FRAMEWORK
  if (environ == NULL)
    environ = *_NSGetEnviron();
#endif
  if (ps->ps_env)
  {
    env = create_environ(ps->ps_env);
    if (env == NULL) {
      PyMem_DEL(argv);
      return NULL;
    }
  } else
    env = NULL;

  /* create the process */
  pid = fork();
  if (pid == -1)
  {
    PyErr_SetString(PyProcess_Error, "unable to fork new process");
    return NULL;
  }
  if (pid == 0)
  {
    PyOS_AfterFork();
    execve(ps->ps_path, argv, env ? env : environ);

    /* if we get here it is definitely an error */
    PySys_WriteStderr("exec of %s failed: %s\n", ps->ps_path, strerror(errno));
    if (env) free_block(env);
    PyMem_DEL(argv);
    exit(1);
  }

  if (env) free_block(env);
  PyMem_DEL(argv);
  ps->ps_pid = pid;
#endif

  return PyInt_FromLong(ps->ps_pid);
}

#ifndef MS_WINDOWS /* we don't need a timer on Windows */

/* these defines will be in the Python.h(config.h) header file */
static double floattime(void)
{
  /* There are three ways to get the time:
    (1) gettimeofday() -- resolution in microseconds
    (2) ftime() -- resolution in milliseconds
    (3) time() -- resolution in seconds
    In all cases the return value is a float in seconds.
    Since on some systems (e.g. SCO ODT 3.0) gettimeofday() may
    fail, so we fall back on ftime() or time().
    Note: clock resolution does not imply clock accuracy! */
#ifdef HAVE_GETTIMEOFDAY
  {
    struct timeval t;
#ifdef GETTIMEOFDAY_NO_TZ
    if (gettimeofday(&t) == 0)
      return (double)t.tv_sec + t.tv_usec*0.000001;
#else /* !GETTIMEOFDAY_NO_TZ */
    if (gettimeofday(&t, (struct timezone *)NULL) == 0)
      return (double)t.tv_sec + t.tv_usec*0.000001;
#endif /* !GETTIMEOFDAY_NO_TZ */
  }
#endif /* !HAVE_GETTIMEOFDAY */
  {
#if defined(HAVE_FTIME)
    struct timeb t;
    ftime(&t);
    return (double)t.time + (double)t.millitm * (double)0.001;
#else /* !HAVE_FTIME */
    time_t secs;
    time(&secs);
    return (double)secs;
#endif /* !HAVE_FTIME */
  }
}


#endif

static PyObject *wait_for_process(PyProcessObject *ps, PyObject *timeout)
{
#ifdef MS_WINDOWS
  DWORD status, milliseconds;

  if (!ps->ps_handle)
  {
    PyErr_SetString(PyProcess_Error, "process not started");
    return NULL;
  }

  if (timeout)
  {
    double temp = PyFloat_AsDouble(timeout);
    milliseconds = (DWORD)(temp * 1000);
  } else {
    milliseconds = INFINITE;
  }

  status = WaitForSingleObject(ps->ps_handle, milliseconds);
  if (status == WAIT_TIMEOUT)
  {
    return Py_BuildValue("(ii)", 0, 0);
  } else if (status == WAIT_OBJECT_0)
  {
    DWORD exitCode;
    if (!GetExitCodeProcess(ps->ps_handle, &exitCode))
    {
      PyErr_SetString(PyProcess_Error, "GetExitCodeProcess failed");
      return NULL;
    }
    return Py_BuildValue("(ii)", 1, exitCode);
  } else
  {
    PyErr_SetString(PyProcess_Error, "WaitForSingleObject failed");
    return NULL;
  }
#else /* not Windows */
  int status, pid = 0;

  if (timeout)
  {
    double end_time = PyFloat_AsDouble(timeout) + floattime();

    do
    {
      if ((pid = waitpid(ps->ps_pid, &status, WNOHANG)) == -1)
      {
        PyErr_SetFromErrno(PyProcess_Error);
        return NULL;
      }
    } while (!pid && end_time <= floattime());
  } else
  {
    /* wait indefinitely */
    if ((pid = waitpid(ps->ps_pid, &status, 0)) == -1)
    {
      PyErr_SetFromErrno(PyProcess_Error);
      return NULL;
    }
  }

  return Py_BuildValue("(ii)",
                         (pid > 0) ? 1 : 0,
                         WIFEXITED(status) ? WEXITSTATUS(status) : 0);
#endif
}

static PyObject *PyProcessObject_New(char *path, PyObject *argv, PyObject *env)
{
  PyProcessObject *self;

  PyProcess_Type.ob_type = &PyType_Type;
  self = PyObject_New(PyProcessObject, &PyProcess_Type);
  if (self == NULL) {
    Py_DECREF(argv);
    Py_XDECREF(env);
    return NULL;
  }

  /* strdup is not ANSI C */
  self->ps_path = PyMem_NEW(char, strlen(path)+1);
  if (self->ps_path == NULL)
  {
    Py_DECREF(argv);
    Py_XDECREF(env);
    return NULL;
  }
  strcpy(self->ps_path, path);
  self->ps_argv = argv;
  self->ps_env = env;
  self->ps_pid = 0;
#ifdef MS_WINDOWS
  self->ps_handle = NULL;
#endif

  return (PyObject *)self;
}

static void process_dealloc(PyProcessObject *self)
{
  PyMem_DEL(self->ps_path);
  Py_DECREF(self->ps_argv);
  Py_XDECREF(self->ps_env);

#ifdef MS_WINDOWS
  if (self->ps_handle != NULL)
  {
    CloseHandle(self->ps_handle);
  }
#endif

  PyObject_Del(self);
}

static char process_start__doc__[] = \
"P.start()\n\
\n\
Start the process.";

static PyObject *process_start(PyObject *self, PyObject *args)
{
  if (!PyArg_ParseTuple(args, ":start"))
    return NULL;

  return create_process((PyProcessObject*)self);
}

static char process_wait__doc__[] = \
"P.wait([timeout]) -> tuple\n\
\n\
Wait for the process to exit, or for the timeout to expire.  Returns\n\
a tuple of exited flag and exit code.  If the exited flag is false, the\n\
exit code is meaningless.";

static PyObject *process_wait(PyObject *self, PyObject *args)
{
  PyObject *timeout = NULL;
  if (!PyArg_ParseTuple(args, "|O:wait", &timeout))
    return NULL;

  if (timeout && !PyNumber_Check(timeout))
  {
    PyErr_SetString(PyExc_TypeError, "wait() arg 1 must be a number");
    return NULL;
  }
  return wait_for_process((PyProcessObject*)self, timeout);
}


static PyMethodDef process_methods[] = {
  {"start",    process_start,    METH_VARARGS, process_start__doc__},
  {"wait",     process_wait,     METH_VARARGS, process_wait__doc__},
  {NULL,    NULL}   /* sentinel */
};

static PyObject *process_getattr(PyObject *self, char *name)
{
  return Py_FindMethod(process_methods, self, name);
}

static int process_setattr(PyObject *self, char *name, PyObject *v)
{
  if (v == NULL) {
    /* remove an attribute from the instance dict */
    PyErr_Format(PyExc_AttributeError,
                 "%.50s instance has no attribute '%.400s'",
                 PyString_AS_STRING(self->ob_type->tp_name), name);
    return -1;
  }
  else
    /* set an attribute in the instance dictionary */
    PyErr_Format(PyExc_TypeError, "object has read-only attributes");
    return -1;
}

/*** the type ***/

static PyTypeObject PyProcess_Type = {
  PyObject_HEAD_INIT(NULL)
  0,      /*ob_size*/
  "process",     /*tp_name*/
  sizeof(PyProcessObject), /*tp_basicsize*/
  0,      /*tp_itemsize*/
  /* methods */
  (destructor)process_dealloc, /*tp_dealloc*/
  0,      /*tp_print*/
  process_getattr, /*tp_getattr*/
  process_setattr, /*tp_setattr*/
  0,      /*tp_compare*/
  0,      /*tp_repr*/
  0,      /*tp_as_number*/
  0,      /*tp_as_sequence*/
  0,      /*tp_as_mapping*/
  0,      /*tp_hash*/
};

/*******************************************************************/
/* Module things start here                                        */
/*******************************************************************/

/*** functions ***/

static char PyProcess_process__doc__[] = \
"Process(path, args[, env]) -> process object\n\
\n\
Create a new process";

static PyObject *PyProcess_process(PyObject *self, PyObject *args)
{
  char *path;
  PyObject *argv;
  PyObject *env = NULL;
  PyObject *key, *val;
  int i, argc;

  if (!PyArg_ParseTuple(args, "sO|O:Process", &path, &argv, &env))
    return NULL;

  if (!(PyList_Check(argv) || PyTuple_Check(argv)))
  {
    PyErr_SetString(PyExc_TypeError, "Process() arg 2 must be a tuple or list");
    return NULL;
  } else if (PySequence_Size(argv) == 0)
  {
    PyErr_SetString(PyExc_TypeError, "Process() arg 2 must not be empty");
    return NULL;
  }

  if (env && !PyMapping_Check(env))
  {
    PyErr_SetString(PyExc_TypeError, "Process() arg 3 must be a mapping object");
    return NULL;
  }

  /* incref existing tuple object or create new tuple from list */
  argc = PySequence_Size(argv);
  argv = PySequence_Tuple(argv);

  /* verify argv sequence */
  for (i = 0; i < argc; i++)
  {
    char *arg;
    if (!PyArg_Parse(PyTuple_GetItem(argv, i),
                     "s;Process() arg 2 must contain only strings",
                     &arg))
    {
      /* error set in PyArg_Parse */
      Py_DECREF(argv);
      return NULL;
    }
  }

  if (env)
  {
    /* verify env mapping */
    int envc = PyMapping_Size(env);
    PyObject *keys = PyMapping_Keys(env);
    PyObject *vals = PyMapping_Values(env);
    if (!keys || !vals)
    {
      /* error already set */
      Py_XDECREF(vals);
      Py_XDECREF(keys);
      Py_DECREF(argv);
      return NULL;
    }

    for (i = 0; i < envc; i++)
    {
      char *k, *v;
      key = PyList_GetItem(keys, i);
      val = PyList_GetItem(vals, i);
      if (!key || !val)
      {
        /* error already set */
        Py_DECREF(vals);
        Py_DECREF(keys);
        Py_DECREF(argv);
        return NULL;
      }

      if (!PyArg_Parse(key, "s;Process() arg 3 contains a non-string key", &k) ||
          !PyArg_Parse(val, "s;Process() arg 3 contains a non-string value", &v))
      {
        /* error already set */
        Py_DECREF(vals);
        Py_DECREF(keys);
        Py_DECREF(argv);
        return NULL;
      }
    }
    Py_DECREF(keys);
    Py_DECREF(vals);
    Py_INCREF(env);
  }
  return PyProcessObject_New(path, argv, env);
}

static char PyProcess_detach__doc__[] = 
"DetachProcess()\n\
\n\
Detaches the calling process for its terminal.";

static PyObject *PyProcess_detach(PyObject *module, PyObject *args)
{
  if (!PyArg_ParseTuple(args, ":DetachProcess"))
    return NULL;

#ifdef MS_WINDOWS
  if (!FreeConsole()) {
    return PyErr_SetFromErrno(PyExc_OSError);
  }
#else /* ndef MS_WINDOWS */

#ifdef HAVE_SETSID
  if (setsid() == -1) {
    return PyErr_SetFromErrno(PyExc_OSError);
  }
#else
  if (setpgid(0, 0) == -1) {
    return PyErr_SetFromErrno(PyExc_OSError);
  }
#endif
  
  /* close out the standard file descriptors */
  if (freopen("/dev/null", "r", stdin) == NULL) {
    return PyErr_SetFromErrno(PyExc_OSError);
  }
  if (freopen("/dev/null", "w", stdout) == NULL) {
    return PyErr_SetFromErrno(PyExc_OSError);
  }

  if (freopen("/dev/null", "w", stderr) == NULL) {
    return PyErr_SetFromErrno(PyExc_OSError);
  }
#endif /* ndef MS_WINDOWS */

  Py_INCREF(Py_None);
  return Py_None;
}

static PyMethodDef PyProcess_methods[] = {
  { "Process", PyProcess_process, METH_VARARGS, PyProcess_process__doc__ },
  { "DetachProcess", PyProcess_detach, METH_VARARGS, PyProcess_detach__doc__ },
  { NULL, NULL }
};

static char PyProcess__doc__[] =
"This module provides functionality for manipulating processes.\n\
\n\
Functions:\n\
\n\
Process() -- create a new process object\n\
\n\
Special objects:\n\
\n\
ProcessType -- type object for process objects\n\
ProcessError -- exception raised for errors\n\
";

DL_EXPORT(void) init_processc(void) {
  PyObject *m, *d;
  PyObject *mod_name;
  char *exc_name;

  m = Py_InitModule3("_processc", PyProcess_methods, PyProcess__doc__);
  d = PyModule_GetDict(m);


  /* make our type a valid Python type */
  PyProcess_Type.ob_type = &PyType_Type;
  Py_INCREF(&PyProcess_Type);

  /* set global error class */
  /* dynamically create the class name based on the module name */
  /* too bad that Python cannot do this for us */
  mod_name = PyDict_GetItemString(d, "__name__"); /* borrowed */
  exc_name = (char *) malloc(PyString_Size(mod_name) + 14);
  sprintf(exc_name, "%s.%s", PyString_AS_STRING(mod_name), "ProcessError");
  PyProcess_Error = PyErr_NewException(exc_name, NULL, NULL);
  free(exc_name);
  if (PyProcess_Error == NULL) {
    printf("IN HERE\n");
    return;
    }

  /* add our globals to the module dict */
  PyDict_SetItemString(d, "ProcessError", PyProcess_Error);
  PyDict_SetItemString(d, "ProcessType", (PyObject *)&PyProcess_Type);
}

