#define _ATFILE_SOURCE
#include <Python.h>
#include "structseq.h"
#include "structmember.h"
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <dirent.h>


/************************************************************/

typedef struct {
  PyObject_HEAD
  int fd;
} dirobject;

static PyTypeObject PyDir_Type;
static PyTypeObject *StatResultType;

#define PyDir_Check(op) PyObject_TypeCheck(op, &PyDir_Type)


static PyObject *_dir_new(PyTypeObject *type, int fd, char *path_for_errors)
{
  if (fd < 0)
    return PyErr_SetFromErrnoWithFilename(PyExc_OSError, path_for_errors);

  dirobject *diro = (dirobject *)type->tp_alloc(type, 0);
  if (diro == NULL)
    {
      close(fd);
      return NULL;
    }
  diro->fd = fd;
  return (PyObject *)diro;
}

static PyObject *dir_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
  static char *kwargs[] = {"dirname", NULL};

  dirobject *diro;
  int fd;
  char *path = ".";
  if (!PyArg_ParseTupleAndKeywords(args, kwds, "|et:dir", kwargs,
                                   Py_FileSystemDefaultEncoding, &path))
    return NULL;

  fd = open(path, O_RDONLY|O_DIRECTORY);
  return _dir_new(type, fd, path);
}

static void dir_dealloc(dirobject *diro)
{
  close(diro->fd);
  PyObject_Del(diro);
}

static int samestat(struct stat *buf1, struct stat *buf2)
{
  return (buf1->st_ino == buf2->st_ino &&
          buf1->st_dev == buf2->st_dev);
}

static void _set_error(int err)
{
  char *s = strerror(err);
  PyObject *v = Py_BuildValue("(is)", err, s);
  if (v != NULL) {
    PyErr_SetObject(PyExc_OSError, v);
    Py_DECREF(v);
  }
}

/* compute the path of fd1 relatively to bufroot.  The buffer must
   be a char[PATH_MAX+1].  If successful, sets a part of the buffer
   to the form "path/path/path" (possibly empty) and returns a pointer
   to it.  Returns NULL with no exception if not found, or with an
   exception set if an error occurred.
*/
static char *relto(int fd, struct stat *bufroot, char *buffer)
{
  char *p = buffer + PATH_MAX;
  *p = '\0';

  struct stat buf;
  if (fstat(fd, &buf) < 0)
    {
      PyErr_SetFromErrno(PyExc_OSError);
      return NULL;
    }
  if (samestat(&buf, bufroot))
    return p;

  int fd1 = openat(fd, "..", O_RDONLY|O_DIRECTORY);

  while (1)
    {
      if (fd1 < 0)
        {
          PyErr_SetFromErrno(PyExc_OSError);
          return NULL;
        }
      DIR *dirp = fdopendir(fd1);
      assert(dirp != NULL);

      while (1)
        {
          errno = 0;
          struct dirent *dirent = readdir(dirp);
          if (dirent == NULL)
            {
              if (errno)
                PyErr_SetFromErrno(PyExc_OSError);
              closedir(dirp);
              return NULL;
            }
          if (dirent->d_type != DT_DIR && dirent->d_type != DT_UNKNOWN)
            continue;

          struct stat subbuf;
          if (fstatat(fd1, dirent->d_name, &subbuf, AT_SYMLINK_NOFOLLOW) < 0)
            continue;	/* xxx ignore fstatat() errors */

          if (samestat(&buf, &subbuf))
            {
              int len = strlen(dirent->d_name);
              p -= len;
              if (p <= buffer)
                {
                  _set_error(ENAMETOOLONG);
                  closedir(dirp);
                  return NULL;
                }
              memcpy(p, dirent->d_name, len);
              p -= 1;
              *p = '/';
              break;
            }
        }

      struct stat buf1;
      if (fstat(fd1, &buf1) < 0)
        {
          PyErr_SetFromErrno(PyExc_OSError);
          closedir(dirp);
          return NULL;
        }
      if (samestat(&buf1, bufroot))
        {
          closedir(dirp);
          break;
        }
      if (samestat(&buf1, &buf))
        {
          /* reached the root, not found */
          return NULL;
        }
      buf = buf1;
      fd1 = openat(fd1, "..", O_RDONLY|O_DIRECTORY);
      closedir(dirp);
    }

  return p + 1;
}

static PyObject *dir_repr(dirobject *diro)
{
  struct stat bufroot;
  if (lstat("/", &bufroot) < 0)
    return PyErr_SetFromErrnoWithFilename(PyExc_OSError, "/");

  char reprbuffer[PATH_MAX+1];
  char *p = relto(diro->fd, &bufroot, reprbuffer);
  if (p == NULL)
    {
      PyErr_Clear();
      return PyString_FromString("<dir '?'>");
    }
  return PyString_FromFormat("<dir '/%s'>", p);
}

static PyObject *dir_relto(dirobject *diro, PyObject *args)
{
  dirobject *parent;
  if (!PyArg_ParseTuple(args, "O!", &PyDir_Type, &parent))
    return NULL;

  struct stat bufroot;
  if (fstat(parent->fd, &bufroot) < 0)
    {
      PyErr_SetFromErrno(PyExc_OSError);
      return NULL;
    }

  char reprbuffer[PATH_MAX+1];
  char *p = relto(diro->fd, &bufroot, reprbuffer);
  if (p == NULL)
    {
      if (PyErr_Occurred())
        return NULL;
      else
        {
          Py_INCREF(Py_None);
          return Py_None;
        }
    }
  if (*p == '\0')
    p = ".";
  return PyString_FromString(p);
}

static PyObject *dir_join(dirobject *diro, PyObject *args)
{
  char *newdir;
  if (!PyArg_ParseTuple(args, "et", Py_FileSystemDefaultEncoding, &newdir))
    return NULL;

  int fd = openat(diro->fd, newdir, O_RDONLY|O_DIRECTORY);
  return _dir_new(&PyDir_Type, fd, newdir);
}

static PyObject *dir_open(dirobject *diro, PyObject *args)
{
  char *filename, *mode = "r";
  int bufsize = -1;
  if (!PyArg_ParseTuple(args, "et|si", Py_FileSystemDefaultEncoding, &filename,
                        &mode, &bufsize))
    return NULL;

  char basemode = *mode == 'U' ? 'r' : *mode;
  int plus = strchr(mode, '+') != NULL;
  int flags;
  switch (basemode) {
  case 'r': flags = plus ? O_RDWR : O_RDONLY; break;
  case 'w': flags = (plus ? O_RDWR : O_WRONLY) | O_CREAT | O_TRUNC; break;
  case 'a': flags = (plus ? O_RDWR : O_WRONLY) | O_CREAT | O_APPEND; break;
  default:
    PyErr_SetString(PyExc_ValueError, "invalid mode");
    return NULL;
  }
  int fd = openat(diro->fd, filename, flags, 0666);
  if (fd < 0)
    return PyErr_SetFromErrnoWithFilename(PyExc_OSError, filename);

  FILE *fp = fdopen(fd, mode);
  if (fp == NULL)
    {
      close(fd);
      return PyErr_SetFromErrnoWithFilename(PyExc_OSError, filename);
    }
  PyObject *f = PyFile_FromFile(fp, filename, mode, fclose);
  if (f != NULL)
    PyFile_SetBufSize(f, bufsize);
  return f;
}

static PyObject *dir_mkdir(dirobject *diro, PyObject *args)
{
  char *dirname;
  if (!PyArg_ParseTuple(args, "et", Py_FileSystemDefaultEncoding, &dirname))
    return NULL;
  if (mkdirat(diro->fd, dirname, 0777) < 0)
    return PyErr_SetFromErrnoWithFilename(PyExc_OSError, dirname);
  Py_INCREF(Py_None);
  return Py_None;
}

static void _fill_time(PyObject *v, int index, time_t sec, unsigned long nsec)
{
  PyObject *fval,*ival;
#if SIZEOF_TIME_T > SIZEOF_LONG
  ival = PyLong_FromLongLong((PY_LONG_LONG)sec);
#else
  ival = PyInt_FromLong((long)sec);
#endif
  if (!ival)
    return;
  fval = PyFloat_FromDouble(sec + 1e-9*nsec);
  PyStructSequence_SET_ITEM(v, index, ival);
  PyStructSequence_SET_ITEM(v, index+3, fval);
}

static PyObject *_do_stat(dirobject *diro, PyObject *args, int flags)
{
  char *filename;
  struct stat st;
  if (!PyArg_ParseTuple(args, "et", Py_FileSystemDefaultEncoding, &filename))
    return NULL;
  if (fstatat(diro->fd, filename, &st, flags) < 0)
    return PyErr_SetFromErrnoWithFilename(PyExc_OSError, filename);

  PyObject *v = PyStructSequence_New(StatResultType);
  if (v == NULL)
    return NULL;

  PyStructSequence_SET_ITEM(v, 0, PyInt_FromLong((long)st.st_mode));
  PyStructSequence_SET_ITEM(v, 1,
                            PyLong_FromLongLong((PY_LONG_LONG)st.st_ino));
  PyStructSequence_SET_ITEM(v, 2,
                            PyLong_FromLongLong((PY_LONG_LONG)st.st_dev));
  PyStructSequence_SET_ITEM(v, 3, PyInt_FromLong((long)st.st_nlink));
  PyStructSequence_SET_ITEM(v, 4, PyInt_FromLong((long)st.st_uid));
  PyStructSequence_SET_ITEM(v, 5, PyInt_FromLong((long)st.st_gid));
  PyStructSequence_SET_ITEM(v, 6,
                            PyLong_FromLongLong((PY_LONG_LONG)st.st_size));

  unsigned long ansec, mnsec, cnsec;
  ansec = st.st_atim.tv_nsec;
  mnsec = st.st_mtim.tv_nsec;
  cnsec = st.st_ctim.tv_nsec;
  _fill_time(v, 7, st.st_atime, ansec);
  _fill_time(v, 8, st.st_mtime, mnsec);
  _fill_time(v, 9, st.st_ctime, cnsec);

  PyStructSequence_SET_ITEM(v, 13,
                            PyInt_FromLong((long)st.st_blksize));
  PyStructSequence_SET_ITEM(v, 14,
                            PyInt_FromLong((long)st.st_blocks));
  PyStructSequence_SET_ITEM(v, 15,
                            PyInt_FromLong((long)st.st_rdev));

  if (PyErr_Occurred()) {
    Py_DECREF(v);
    return NULL;
  }
  return v;
}

static PyObject *dir_stat(dirobject *diro, PyObject *args)
{
  return _do_stat(diro, args, 0);
}

static PyObject *dir_lstat(dirobject *diro, PyObject *args)
{
  return _do_stat(diro, args, AT_SYMLINK_NOFOLLOW);
}

static PyObject *dir_symlink(dirobject *diro, PyObject *args)
{
  char *oldname, *newname;
  if (!PyArg_ParseTuple(args, "etet", Py_FileSystemDefaultEncoding, &oldname,
                                      Py_FileSystemDefaultEncoding, &newname))
    return NULL;
  if (symlinkat(oldname, diro->fd, newname) < 0)
    return PyErr_SetFromErrnoWithFilename(PyExc_OSError, newname);
  Py_INCREF(Py_None);
  return Py_None;
}

static PyObject *dir_rename(dirobject *diro, PyObject *args)
{
  dirobject *newdiro;
  char *oldname, *newname;
  if (!PyArg_ParseTuple(args, "etO!et",
                        Py_FileSystemDefaultEncoding, &oldname,
                        &PyDir_Type, &newdiro,
                        Py_FileSystemDefaultEncoding, &newname))
    return NULL;
  if (renameat(diro->fd, oldname, newdiro->fd, newname) < 0)
    {
      PyErr_SetFromErrno(PyExc_OSError);
      return NULL;
    }
  Py_INCREF(Py_None);
  return Py_None;
}

static PyObject *_do_unlink(dirobject *diro, PyObject *args, int flags)
{
  char *filename;
  if (!PyArg_ParseTuple(args, "et", Py_FileSystemDefaultEncoding, &filename))
    return NULL;
  if (unlinkat(diro->fd, filename, flags) < 0)
    return PyErr_SetFromErrnoWithFilename(PyExc_OSError, filename);
  Py_INCREF(Py_None);
  return Py_None;
}

static PyObject *dir_unlink(dirobject *diro, PyObject *args)
{
  return _do_unlink(diro, args, 0);
}

static PyObject *dir_rmdir(dirobject *diro, PyObject *args)
{
  return _do_unlink(diro, args, AT_REMOVEDIR);
}

static PyObject *dir_readlink(dirobject *diro, PyObject *args)
{
  char buf[PATH_MAX];
  char *filename;
  if (!PyArg_ParseTuple(args, "et", Py_FileSystemDefaultEncoding, &filename))
    return NULL;
  int n = readlinkat(diro->fd, filename, buf, sizeof(buf));
  if (n < 0)
    return PyErr_SetFromErrnoWithFilename(PyExc_OSError, filename);
  return PyString_FromStringAndSize(buf, n);
}

static PyObject *dir_link(dirobject *diro, PyObject *args)
{
  dirobject *newdiro;
  char *oldname, *newname;
  if (!PyArg_ParseTuple(args, "etO!et",
                        Py_FileSystemDefaultEncoding, &oldname,
                        &PyDir_Type, &newdiro,
                        Py_FileSystemDefaultEncoding, &newname))
    return NULL;
  if (linkat(diro->fd, oldname, newdiro->fd, newname, 0) < 0)
    {
      PyErr_SetFromErrno(PyExc_OSError);
      return NULL;
    }
  Py_INCREF(Py_None);
  return Py_None;
}

static PyMethodDef Dir_methods[] = {
  {"join",	  (PyCFunction)dir_join, METH_VARARGS},
  {"relto",	  (PyCFunction)dir_relto, METH_VARARGS},
  {"open",	  (PyCFunction)dir_open, METH_VARARGS},
  {"mkdir",	  (PyCFunction)dir_mkdir, METH_VARARGS},
  {"stat",	  (PyCFunction)dir_stat, METH_VARARGS},
  {"lstat",	  (PyCFunction)dir_lstat, METH_VARARGS},
  {"symlink",	  (PyCFunction)dir_symlink, METH_VARARGS},
  {"rename",	  (PyCFunction)dir_rename, METH_VARARGS},
  {"unlink",	  (PyCFunction)dir_unlink, METH_VARARGS},
  {"rmdir",	  (PyCFunction)dir_rmdir, METH_VARARGS},
  {"readlink",	  (PyCFunction)dir_readlink, METH_VARARGS},
  {"link",	  (PyCFunction)dir_link, METH_VARARGS},
  {NULL,	  NULL}		/* sentinel */
};

static PyMemberDef Dir_members[] = {
  {"_fileno",  T_INT, offsetof(dirobject, fd),  READONLY},
  {NULL}
};

static PyTypeObject PyDir_Type = {
  PyObject_HEAD_INIT(NULL)
  0,				/* ob_size */
  "pathat.dir",		/* tp_name */
  sizeof(dirobject),		/* tp_basicsize */
  0,				/* tp_itemsize */
  /* methods */
  (destructor)dir_dealloc,	/* tp_dealloc */
  0,				/* tp_print */
  0,				/* tp_getattr */
  0,				/* tp_setattr */
  0,				/* tp_compare */
  (unaryfunc)dir_repr,		/* tp_repr */
  0,				/* tp_as_number */
  0,				/* tp_as_sequence */
  0,				/* tp_as_mapping */
  0,				/* tp_hash */
  0,				/* tp_call */
  0,				/* tp_str */
  PyObject_GenericGetAttr,	/* tp_getattro */
  0,				/* tp_setattro */
  0,				/* tp_as_buffer */
  Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,	/* tp_flags */
  0,				/* tp_doc */
  0,				/* tp_traverse */
  0,				/* tp_clear */
  0,				/* tp_richcompare */
  0,				/* tp_weaklistoffset */
  0,				/* tp_iter */
  0,				/* tp_iternext */
  Dir_methods,			/* tp_methods */
  Dir_members,			/* tp_members */
  0,				/* tp_getset */
  0,				/* tp_base */
  0,				/* tp_dict */
  0,				/* tp_descr_get */
  0,				/* tp_descr_set */
  0,				/* tp_dictoffset */
  0,				/* tp_init */
  0,				/* tp_alloc */
  dir_new,			/* tp_new */
  PyObject_GC_Del,		/* tp_free */
};

/************************************************************/

PyDoc_STRVAR(module_doc,
"Access to the Linux-specific functions openat & co.\n\
");

static PyMethodDef module_methods[] = {
 	{NULL,		NULL}		/* sentinel */
};

PyMODINIT_FUNC
initpathat(void)
{
  PyObject *m = Py_InitModule3("pathat", module_methods, module_doc);
  if (m == NULL)
    return;

  PyObject* os_module = PyImport_ImportModule("posix");
  if (os_module == NULL)
    return;
  PyObject* sr = PyObject_GetAttrString(os_module, "stat_result");
  if (sr == NULL)
    return;
  if (!PyType_Check(sr))
    {
      PyErr_SetString(PyExc_TypeError, "posix.stat_result: type expected");
      return;
    }
  StatResultType = (PyTypeObject *)sr;

  if (PyType_Ready(&PyDir_Type) < 0)
    return;
  PyModule_AddObject(m, "dir", (PyObject *)&PyDir_Type);
}