#include "precompile.hxx"

#define LIBNAME   executeAsync
#define STR(s) #s
#define LIBSNAME STR(LIBNAME)
#define MAKELUAOPEN(name) luaopen_ ## name
#define LUAOPEN(name) MAKELUAOPEN(name)

static const char* metatableKey = LIBSNAME" Metatable Key";

int pushError (lua_State* L) {
  int err = errno; // calls to Lua may modify errno
  lua_pushnil (L);
  lua_pushstring (L, strerror (err));
  return 2;
}

#define nil NULL
struct FdsAndId {std::array<int, 2> fds; int id;};
class Jobs {
  fd_set readfds;

  std::map <pid_t, FdsAndId> PidMap;
public:
  Jobs () {
    ++jobQueueCount;
    miId = jobQueueCount;
  }

  int execute (lua_State* L, const char* cmdline, int id) {
    int pipeOut [2];
    int pipeErr [2];
    pid_t pidChild;
    if (pipe2 (pipeOut, O_NONBLOCK) == -1) return pushError (L);
    if (pipe2 (pipeErr, O_NONBLOCK) == -1) return pushError (L);
    pidChild = fork();
    if (pidChild == -1) return pushError (L);

    constexpr int readEnd = 0;
    constexpr int writeEnd = 1;
    if (pidChild == 0) {    /* Child writes to pipe */
      close (pipeOut [readEnd]);
      dup2 (pipeOut [writeEnd], STDOUT_FILENO);
      // Check `man stdout` for more info
      close (pipeErr [readEnd]);
      dup2 (pipeErr [writeEnd], STDERR_FILENO);
      // Check `man stdout` for more info
      close (pipeOut [writeEnd]);
      close (pipeErr [writeEnd]);
      execlp ("sh", "sh", "-c", cmdline, (char*) NULL);
      return 0;
    } else { /* Parent */
      close (pipeOut [writeEnd]);
      close (pipeErr [writeEnd]);
      PidMap [pidChild] = {{pipeOut[readEnd], pipeErr[readEnd]}, id};
      lua_pushboolean (L, 1);
      lua_pushinteger (L, pidChild);
      lua_pushinteger (L, PidMap.size ());
      return 3;
    }
  } // end Jobs::execute

#define addMainReturn(index) lua_rawseti (L, -2, index)
#define addSubReturn(index) lua_rawseti (L, -2, index)

  int wait (lua_State* L) {
    if (! PidMap.size ()) {
      lua_createtable (L, 1, 0); // main return
      lua_createtable (L, 3, 0); // sub return
      lua_pushinteger (L, 0); addSubReturn (1);
      lua_pushinteger (L, -1); addSubReturn (2);
      lua_pushinteger (L, -1); addSubReturn (3);
      addMainReturn (1);
      return 1; // lua: {{0, -1, -1}}
    }

    /* Initialize readfds, before the select() call. */
    FD_ZERO (&readfds);
    int maxfds = 0;
    for (const auto &i: PidMap) {
      for (const auto& j: i.second.fds) {
        FD_SET (j, &readfds);
        maxfds = std::max (maxfds, j);
      }
    }
    int r = select (maxfds+1, &readfds, NULL, nil, NULL);
    if (r == -1) {
      /* Handle error */
      const char* err = 0;
      switch (errno) {
      case EBADF:
        err = "One or more of the file descriptor sets specified a file"
          " descriptor that is not a valid open file descriptor.";
        break;
      case EINTR:
        err = "The function was interrupted while blocked waiting for any"
          " of the selected descriptors to become ready and before the"
          " timeout interval expired.";
        break;
      case EINVAL:
        err = "The nfds argument is less than 0 or greater than FD_SETSIZE."
          " OR One of the specified file descriptors refers to a STREAM"
          " or multiplexer that is linked (directly or indirectly)"
          " downstream from a multiplexer.";
        break;
      }
      lua_pushnil (L);
      lua_pushstring (L, err);
      return 2;
    } // end error

    lua_createtable (L, r, 0); // main return
    int indexInReturnTable = 1;
    int countSet = 0;
    for (auto ii = PidMap.begin (); ii != PidMap.end ();) {
      auto& i = *ii;
      auto next = ii;
      ++next;
      const std::array<int, 2>& fds = i.second.fds;
      for (int j = 0; j < fds.size (); ++j) {
        int fd = fds [j];
        if (FD_ISSET (fd, &readfds)) {
          ++countSet;
          FILE* fp = fdopen (fd, "r");
          size_t n = ~((size_t)0);
          size_t rlen = LUAL_BUFFERSIZE;  /* try to read that much each time */
          size_t nr;  /* number of chars actually read */
          luaL_Buffer b;
          luaL_buffinit(L, &b);
          do {
            char *p = luaL_prepbuffer(&b);
            if (rlen > n) rlen = n;  /* cannot read more than asked */
            nr = fread(p, sizeof(char), rlen, fp);
            luaL_addsize(&b, nr);
            n -= nr;  /* still have to read `n' chars */
          } while (n > 0 && nr == rlen);  /* until end of count or eof */
          luaL_pushresult(&b);  /* close buffer */

          int id = i.second.id;
          if (lua_objlen(L, -1) > 0) { // no point in returning empty string
            lua_createtable (L, 3, 0); // sub return
            lua_insert (L, -2); // buffered string comes to the top
            addSubReturn (1);
            lua_pushinteger (L, id);  addSubReturn (2);
            lua_pushinteger (L, j+1); addSubReturn (3);
            addMainReturn (indexInReturnTable++);
          } else {
            lua_pop (L, 1); // the empty string
          }

          if (nr < rlen) {
            if (feof (fp) && j == 0) { // child closed its stdout
              /* this is kinda wrong, should wait for SIGCHILD instead, but will do
               for now */
              PidMap.erase (ii);
              lua_createtable (L, 3, 0); // sub return
              lua_pushinteger (L, PidMap.size ()); addSubReturn (1);
              lua_pushinteger (L, id); addSubReturn (2);
              int exitCode;
              waitpid (i.first, &exitCode, 0);
              lua_pushinteger (L, exitCode); addSubReturn (3);
              addMainReturn (indexInReturnTable++);
            }
          }
        }
      } // end for (int j = 0; j < fds.size (); ++j)
      ii = next;
    } // end for (auto ii = PidMap.begin (); ii != PidMap.end ();)

    if (countSet) return 1; // main return
    else {
      lua_pop (L, 1); // pop main return table
      lua_pushnil (L);
      lua_pushstring (L, "select triggered but no fd is FD_ISSET!");
      return 2;
    }
  } // end ::wait
private:
  int miId;
  static int jobQueueCount;
};
int Jobs::jobQueueCount = 0;

Jobs *check_self (lua_State *L, int index = 1) {
  Jobs *p = (Jobs *) lua_touserdata(L, index);
  if (p != NULL) {  /* value is a userdata? */
    if (lua_getmetatable(L, index)) {  /* does it have a metatable? */
      /* get supposed metatable */
      lua_pushlightuserdata (L, &metatableKey);
      lua_rawget (L, LUA_REGISTRYINDEX);
      if (lua_rawequal(L, -1, -2)) {  /* does it have the correct mt? */
        lua_pop(L, 2);  /* remove both metatables */
        return p;
      }
    }
  }
  luaL_typerror(L, index, LIBSNAME);  /* else error */
  return NULL;  /* to avoid warnings */
}

extern "C" {

  static int new_jobs (lua_State *L) {
    new (lua_newuserdata (L, sizeof (Jobs))) Jobs ();

    lua_pushlightuserdata(L, &metatableKey);
    lua_rawget(L, LUA_REGISTRYINDEX);
    lua_setmetatable (L, -2);

    return 1;
  }

  static int execute (lua_State *L) {
    Jobs *self = check_self (L);

    // app: string, args: string, id: int
    const char* cmdline = luaL_checkstring (L, 2);
    int id = luaL_checkint (L, 3);
    return self->execute (L, cmdline, id);
  }

  // returns number of remaining process, id of finished process, exit code of
  // finished process, read handle of output of finished process
  static int waitJob (lua_State *L) {
    Jobs *self = check_self (L);

    return self->wait (L);
  }

  static int gc (lua_State *L) { // '__gc'
    Jobs* self = check_self (L);
    self->~Jobs ();
    return 0;
  }

  static const struct luaL_Reg Jobs_methods [] = {
    {"__gc", gc},
    {NULL, NULL},
  };

  EXPORT int LUAOPEN(LIBNAME) (lua_State *L) {
    lua_pushlightuserdata (L, &metatableKey); // key-in-registry
    lua_newtable (L); // metatable-userdata
    luaL_register (L, NULL, Jobs_methods);

    const struct luaL_Reg Jobs_functions [] = {
      {"execute",  execute },
      {"wait",     waitJob },
      {NULL,       NULL},
    };
    lua_newtable (L); // this will be __index
    luaL_register (L, NULL, Jobs_functions);
    lua_setfield (L, -2, "__index");

    // registry [key-in-registry] = metatable-userdata
    lua_settable (L, LUA_REGISTRYINDEX);

    lua_newtable (L); // class table, to return

    // metatable for class table
    // allow call syntax:
    lua_newtable (L);
    lua_pushcfunction (L, new_jobs);
    lua_setfield (L, -2, "__call");

    lua_setmetatable (L, -2);

    return 1;
  }
} // extern "C"