/*

from distutils.core import setup, Extension

setup(
    name='winio',
    ext_modules=[Extension('winfileio', ['winfileio.c'])],
    )

*/

#include "Python.h"
#include "pythread.h"

#define WINDOWS_LEAN_AND_MEAN
#include <windows.h>

/*
 * Convert a Python integer to a HANDLE (for use with PyArg_ParseTuple())
 */

static int
HandleConverter(PyObject *obj, HANDLE *h)
{
    size_t n = PyLong_AsSize_t(obj);
    if (n == (size_t)-1 && PyErr_Occurred())
        return 0;
    *h = (HANDLE)n;
    return 1;
}

#define CHECK_NOT_CLOSED(f)                                     \
    if ((f)->handle == INVALID_HANDLE_VALUE) {                  \
        PyErr_SetString(PyExc_ValueError, "file closed");       \
        return NULL;                                            \
    }

#define CHECK_READABLE(f)                                       \
    CHECK_NOT_CLOSED(f)                                         \
    if (!(f)->readable) {                                       \
        PyErr_SetString(PyExc_ValueError, "file not readable"); \
        return NULL;                                            \
    }

#define CHECK_WRITABLE(f)                                       \
    CHECK_NOT_CLOSED(f)                                         \
    if (!(f)->writable) {                                       \
        PyErr_SetString(PyExc_ValueError, "file not writable"); \
        return NULL;                                            \
    }

#if SIZEOF_SIZE_T > SIZEOF_INT
# define CLIP_SSIZE_T(n) Py_MIN((n), (Py_ssize_t)UINT_MAX)
#else
# define CLIP_SSIZE_T(n) (n)
#endif

#define MODE(f) mode_string[(f)->readable + (f)->writable*2 - 1]

static const char *mode_string[] = {"r", "w", "r+"};

typedef struct _winio {
    struct {
        PyObject_HEAD
        PyObject *dict;
        PyObject *weakreflist;
    } base;
    HANDLE handle;
    unsigned readable : 1;
    unsigned writable : 1;
    unsigned closehandle : 1;
    // XXX append : 1;
} winio;

static int
winio_init(winio *self, PyObject *args, PyObject *kwds)
{
    HANDLE handle;
    char *mode = "r";
    int closehandle = TRUE;
    static char *kwlist[] = {"handle", "mode", "closehandle", NULL};

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&|si", kwlist,
                                     HandleConverter, &handle,
                                     &mode, &closehandle))
        return -1;
    if (handle == INVALID_HANDLE_VALUE || handle == NULL) {
        PyErr_SetString(PyExc_ValueError, "invalid handle value");
        return -1;
    }
    if (self->handle != NULL) {
        PyErr_SetString(PyExc_ValueError, "cannot reinitialize");
        return -1;
    }
    if (strcmp(mode, "r") == 0) {
        self->readable = 1;
    } else if (strcmp(mode, "w") == 0) {
        self->writable = 1;
    } else if (strcmp(mode, "r+") == 0 || strcmp(mode, "w+") == 0) {
        self->readable = self->writable = 1;
    } else {
        PyErr_Format(PyExc_ValueError, "invalid mode: '%s'; "
                     "expected 'r', 'w', 'r+' or 'w+'", mode);
        return -1;
    }
    if (closehandle)
        self->closehandle = 1;
    self->handle = handle;
    return 0;
}

static PyObject *
winio_close(winio *self)
{
    HANDLE h = self->handle;
    BOOL success;
    if (h != INVALID_HANDLE_VALUE && h != NULL) {
        self->handle = INVALID_HANDLE_VALUE;
        if (self->closehandle) {
            Py_BEGIN_ALLOW_THREADS
            success = CloseHandle(h);
            Py_END_ALLOW_THREADS
            if (!success)
                PyErr_SetFromWindowsErr(0);
        }
    }
    Py_RETURN_NONE;
}

static PyObject *
winio_read(winio *self, PyObject *args)
{
    BOOL success;
    Py_ssize_t size = -1;
    DWORD nbytes;
    PyObject *buf;

    CHECK_READABLE(self);

    if (!PyArg_ParseTuple(args, "|n", &size))
        return NULL;

    if (size < 0) {
        _Py_IDENTIFIER(readall);
        return _PyObject_CallMethodId((PyObject*)self, &PyId_readall, "");
    }
    size = CLIP_SSIZE_T(size);

    buf = PyBytes_FromStringAndSize(NULL, size);
    if (buf == NULL || size == 0)
        return buf;

    Py_BEGIN_ALLOW_THREADS
    success = ReadFile(self->handle, PyBytes_AS_STRING(buf),
                       size, &nbytes, FALSE);
    Py_END_ALLOW_THREADS

    switch (success ? ERROR_SUCCESS : GetLastError()) {
        case ERROR_BROKEN_PIPE:
            nbytes = 0;
            /* fall through */
        case ERROR_SUCCESS:
            if (nbytes == size || _PyBytes_Resize(&buf, nbytes) >= 0)
                return buf;
            break;
        default:
            PyErr_SetFromWindowsErr(0);
    }

    Py_DECREF(buf);
    return NULL;
}

static PyObject *
winio_readinto(winio *self, PyObject *obj)
{
    BOOL success;
    DWORD nbytes;
    DWORD size;
    Py_buffer buf;

    CHECK_READABLE(self);

    if (PyObject_GetBuffer(obj, &buf, PyBUF_WRITABLE) < 0)
        return NULL;

    size = CLIP_SSIZE_T(buf.len);
    if (size == 0) {
        PyBuffer_Release(&buf);
        return PyLong_FromLong(0);
    }

    Py_BEGIN_ALLOW_THREADS
    success = ReadFile(self->handle, buf.buf, size, &nbytes, NULL);
    Py_END_ALLOW_THREADS

    PyBuffer_Release(&buf);

    switch (success ? ERROR_SUCCESS : GetLastError()) {
        case ERROR_SUCCESS:
            return PyLong_FromUnsignedLong(nbytes);
        case ERROR_BROKEN_PIPE:
            return PyLong_FromLong(0);
        default:
            return PyErr_SetFromWindowsErr(0);
    }
}

static PyObject *
winio_write(winio *self, PyObject *obj)
{
    DWORD size;
    DWORD nbytes;
    BOOL success;
    Py_buffer buf;

    CHECK_WRITABLE(self);

    if (PyObject_GetBuffer(obj, &buf, PyBUF_SIMPLE) < 0)
        return NULL;

    size = CLIP_SSIZE_T(buf.len);

    Py_BEGIN_ALLOW_THREADS
    success = WriteFile(self->handle, buf.buf, size, &nbytes, NULL);
    Py_END_ALLOW_THREADS

    PyBuffer_Release(&buf);

    if (success)
        return PyLong_FromUnsignedLong(nbytes);
    else
        return PyErr_SetFromWindowsErr(0);
}

static PyObject *
winio_tell(winio *self)
{
    LARGE_INTEGER offset, newpos;
    BOOL success;

    offset.QuadPart = 0;

    Py_BEGIN_ALLOW_THREADS
    success = SetFilePointerEx(self->handle, offset, &newpos, FILE_CURRENT);
    Py_END_ALLOW_THREADS

    if (!success)
        return PyErr_SetFromWindowsErr(0);

    return PyLong_FromLongLong(newpos.QuadPart);
}

static PyObject *
winio_seek(winio *self, PyObject *args)
{
    LARGE_INTEGER offset, newpos;
    DWORD whence = FILE_BEGIN;
    BOOL success;

    if (!PyArg_ParseTuple(args, "L|k", &offset.QuadPart, &whence))
        return NULL;

    Py_BEGIN_ALLOW_THREADS
    success = SetFilePointerEx(self->handle, offset, &newpos, whence);
    Py_END_ALLOW_THREADS

    if (!success)
        return PyErr_SetFromWindowsErr(0);

    return PyLong_FromLongLong(newpos.QuadPart);
}

static PyObject *
winio_truncate(winio *self, PyObject *args)
{
    LARGE_INTEGER zero, size, oldsize;
    BOOL success;

    CHECK_WRITABLE(self);

    if (!PyArg_ParseTuple(args, "L", &size.QuadPart))
        return NULL;

    if (size.QuadPart < 0) {
        PyErr_SetString(PyExc_ValueError, "negative size");
        return NULL;
    }

    zero.QuadPart = 0;

    Py_BEGIN_ALLOW_THREADS
    /* XXX there is a thread race here -- msvcrt uses a lock */
    success = (SetFilePointerEx(self->handle, zero, &oldsize, FILE_CURRENT) &&
               SetFilePointerEx(self->handle, size, NULL, FILE_BEGIN) &&
               SetEndOfFile(self->handle));
    if (success && oldsize.QuadPart < size.QuadPart)
        success = SetFilePointerEx(self->handle, oldsize, NULL, FILE_BEGIN);
    Py_END_ALLOW_THREADS

    if (!success)
        return PyErr_SetFromWindowsErr(0);

    Py_RETURN_NONE;
}

static PyObject *
winio_repr(winio *self)
{
    return PyUnicode_FromFormat(
        "%s(handle=%zu, mode='%s')",
        Py_TYPE(self)->tp_name, self->handle,
        MODE(self));
}

static PyObject *
winio_readable(winio *self)
{
    return PyBool_FromLong(self->readable);
}

static PyObject *
winio_writable(winio *self)
{
    return PyBool_FromLong(self->writable);
}

static PyObject *
winio_seekable(winio *self)
{
    Py_RETURN_TRUE;
}

static PyObject *
winio_gethandle(winio *self)
{
    if (self->handle == INVALID_HANDLE_VALUE)
        Py_RETURN_NONE;
    else
        return PyLong_FromSize_t((size_t)self->handle);
}

static PyObject *
winio_getmode(winio *self, void *closure)
{
    return PyUnicode_InternFromString(MODE(self));
}

static PyObject *
winio_getclosed(winio *self)
{
    return PyBool_FromLong(self->handle == INVALID_HANDLE_VALUE);
}

static PyMethodDef winio_methods[] = {
    {"read",     (PyCFunction)winio_read,      METH_VARARGS},
    {"readinto", (PyCFunction)winio_readinto,  METH_O,     },
    {"write",    (PyCFunction)winio_write,     METH_O,     },
    {"close",    (PyCFunction)winio_close,     METH_NOARGS },
    {"tell",     (PyCFunction)winio_tell,      METH_NOARGS },
    {"seek",     (PyCFunction)winio_seek,      METH_VARARGS},
    {"truncate", (PyCFunction)winio_truncate,  METH_VARARGS},
    {"readable", (PyCFunction)winio_readable,  METH_NOARGS },
    {"writable", (PyCFunction)winio_writable,  METH_NOARGS },
    {"seekable", (PyCFunction)winio_seekable,  METH_NOARGS },
    {NULL}
};

static PyGetSetDef winio_getsets[] = {
    {"handle",      (getter)winio_gethandle      },
    {"mode",        (getter)winio_getmode        },
    {"closed",      (getter)winio_getclosed      },
    {NULL}
};

PyTypeObject winio_Type = {
    PyVarObject_HEAD_INIT(NULL, 0)
    "winio.winio",                             /* tp_name */
    sizeof(winio),                             /* tp_basicsize */
    0,                                         /* tp_itemsize */
    0,                                         /* tp_dealloc */
    0,                                         /* tp_print */
    0,                                         /* tp_getattr */
    0,                                         /* tp_setattr */
    0,                                         /* tp_reserved */
    (reprfunc)winio_repr,                      /* tp_repr */
    0,                                         /* tp_as_number */
    0,                                         /* tp_as_sequence */
    0,                                         /* tp_as_mapping */
    0,                                         /* tp_hash */
    0,                                         /* tp_call */
    0,                                         /* tp_str */
    0,                                         /* 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 */
    winio_methods,                             /* tp_methods */
    0,                                         /* tp_members */
    winio_getsets,                             /* tp_getset */
    0,                                         /* tp_base */
    0,                                         /* tp_dict */
    0,                                         /* tp_descr_get */
    0,                                         /* tp_descr_set */
    0,                                         /* tp_dictoffset */
    (initproc)winio_init,                      /* tp_init */
};

static struct PyModuleDef winio_module = {
    PyModuleDef_HEAD_INIT,
    "winio", NULL, -1, NULL
};

PyMODINIT_FUNC
PyInit_winfileio(void)
{
    PyObject *m, *_io, *_PyRawIOBase_Type;

    _io = PyImport_ImportModule("_io");
    if (_io == NULL)
        return NULL;

    _PyRawIOBase_Type = PyObject_GetAttrString(_io, "_RawIOBase");
    Py_DECREF(_io);
    if (_PyRawIOBase_Type == NULL)
        return NULL;

    winio_Type.tp_base = (PyTypeObject*)_PyRawIOBase_Type;
    if (PyType_Ready(&winio_Type) < 0)
        return NULL;

    m = PyModule_Create(&winio_module);
    if (m == NULL)
        return NULL;

    Py_INCREF(&winio_Type);
    PyModule_AddObject(m, "WinFileIO", (PyObject*)&winio_Type);
    return m;
}