#include "Python.h"
#include <stdlib.h>
#include <limits.h>

/* datetime.datetime.now() */

/*
    This one is the only function using 'vectorcall' on Python 3.8.
*/

static PyObject*
#if PY_VERSION_HEX >= 0x03070000
_time_machine_now(PyTypeObject *type, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
#else
_time_machine_now(PyTypeObject *type, PyObject **args, Py_ssize_t nargs, PyObject *kwnames)
#endif
{
    PyObject *result = NULL;

    PyObject *time_machine_module = PyImport_ImportModule("time_machine");
    PyObject *time_machine_now = PyObject_GetAttrString(time_machine_module, "now");

#if PY_VERSION_HEX >= 0x03080000
    result = _PyObject_Vectorcall(time_machine_now, args, nargs, kwnames);
#else
    /*
        Parse the given arguments into tz - copied from
        datetime_datetime_new code generated by argument clinic
    */
    static const char * const _keywords[] = {"tz", NULL};
    static _PyArg_Parser _parser = {"|O:now", _keywords, 0};
    PyObject *tz = Py_None;

#if PY_VERSION_HEX >= 0x03070000
    if (!_PyArg_ParseStackAndKeywords(args, nargs, kwnames, &_parser, &tz)) {
        goto exit;
    }
#else
    if (!_PyArg_ParseStack(args, nargs, kwnames, &_parser, &tz)) {
        goto exit;
    }
#endif

    /* Call with a new tuple */
    PyObject* now_args = PyTuple_Pack(1, tz);
    result = PyObject_CallObject(time_machine_now, now_args);
    Py_DECREF(now_args);

exit:
#endif

    Py_DECREF(time_machine_now);
    Py_DECREF(time_machine_module);

    return result;
}

#if PY_VERSION_HEX >= 0x03070000
_PyCFunctionFastWithKeywords original_now = NULL;
#else
_PyCFunctionFast original_now = NULL;
#endif

static PyObject*
#if PY_VERSION_HEX >= 0x03070000
_time_machine_original_now(PyTypeObject *type, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
#else
_time_machine_original_now(PyTypeObject *type, PyObject **args, Py_ssize_t nargs, PyObject *kwnames)
#endif
{
    PyObject *datetime_module = PyImport_ImportModule("datetime");
    PyObject *datetime_class = PyObject_GetAttrString(datetime_module, "datetime");

    PyObject* result = original_now(datetime_class, args, nargs, kwnames);

    Py_DECREF(datetime_class);
    Py_DECREF(datetime_module);

    return result;
}
PyDoc_STRVAR(original_now_doc,
"original_now() -> datetime\n\
\n\
Call datetime.datetime.now() after patching.");

/* datetime.datetime.utcnow() */

static PyObject*
_time_machine_utcnow(PyObject *cls, PyObject *args)
{
    PyObject *time_machine_module = PyImport_ImportModule("time_machine");
    PyObject *time_machine_utcnow = PyObject_GetAttrString(time_machine_module, "utcnow");

    PyObject* result = PyObject_CallObject(time_machine_utcnow, args);

    Py_DECREF(time_machine_utcnow);
    Py_DECREF(time_machine_module);

    return result;
}

PyCFunction original_utcnow = NULL;

static PyObject*
_time_machine_original_utcnow(PyObject *cls, PyObject *args)
{
    PyObject *datetime_module = PyImport_ImportModule("datetime");
    PyObject *datetime_class = PyObject_GetAttrString(datetime_module, "datetime");

    PyObject* result = original_utcnow(datetime_class, args);

    Py_DECREF(datetime_class);
    Py_DECREF(datetime_module);

    return result;
}
PyDoc_STRVAR(original_utcnow_doc,
"original_utcnow() -> datetime\n\
\n\
Call datetime.datetime.utcnow() after patching.");

/* time.time() */

static PyObject*
_time_machine_time(PyObject *self, PyObject *args)
{
    PyObject *time_machine_module = PyImport_ImportModule("time_machine");
    PyObject *time_machine_time = PyObject_GetAttrString(time_machine_module, "time");

    PyObject* result = PyObject_CallObject(time_machine_time, args);

    Py_DECREF(time_machine_time);
    Py_DECREF(time_machine_module);

    return result;
}

PyCFunction original_time = NULL;

static PyObject*
_time_machine_original_time(PyObject *self, PyObject *args)
{
    return original_time(self, args);
}
PyDoc_STRVAR(original_time_doc,
"original_time() -> floating point number\n\
\n\
Call time.time() after patching.");

/* time.localtime() */

static PyObject*
_time_machine_localtime(PyObject *self, PyObject *args)
{
    PyObject *time_machine_module = PyImport_ImportModule("time_machine");
    PyObject *time_machine_localtime = PyObject_GetAttrString(time_machine_module, "localtime");

    PyObject* result = PyObject_CallObject(time_machine_localtime, args);

    Py_DECREF(time_machine_localtime);
    Py_DECREF(time_machine_module);

    return result;
}

PyCFunction original_localtime = NULL;

static PyObject*
_time_machine_original_localtime(PyObject *self, PyObject *args)
{
    return original_localtime(self, args);
}
PyDoc_STRVAR(original_localtime_doc,
"original_localtime() -> floating point number\n\
\n\
Call time.localtime() after patching.");

/* time.gmtime() */

static PyObject*
_time_machine_gmtime(PyObject *self, PyObject *args)
{
    PyObject *time_machine_module = PyImport_ImportModule("time_machine");
    PyObject *time_machine_gmtime = PyObject_GetAttrString(time_machine_module, "gmtime");

    PyObject* result = PyObject_CallObject(time_machine_gmtime, args);

    Py_DECREF(time_machine_gmtime);
    Py_DECREF(time_machine_module);

    return result;
}

PyCFunction original_gmtime = NULL;

static PyObject*
_time_machine_original_gmtime(PyObject *self, PyObject *args)
{
    return original_gmtime(self, args);
}
PyDoc_STRVAR(original_gmtime_doc,
"original_gmtime() -> floating point number\n\
\n\
Call time.gmtime() after patching.");

/* time.strftime() */

static PyObject*
_time_machine_strftime(PyObject *self, PyObject *args)
{
    PyObject *time_machine_module = PyImport_ImportModule("time_machine");
    PyObject *time_machine_strftime = PyObject_GetAttrString(time_machine_module, "strftime");

    PyObject* result = PyObject_CallObject(time_machine_strftime, args);

    Py_DECREF(time_machine_strftime);
    Py_DECREF(time_machine_module);

    return result;
}

PyCFunction original_strftime = NULL;

static PyObject*
_time_machine_original_strftime(PyObject *self, PyObject *args)
{
    return original_strftime(self, args);
}
PyDoc_STRVAR(original_strftime_doc,
"original_strftime() -> floating point number\n\
\n\
Call time.strftime() after patching.");


static PyObject*
_time_machine_patch(PyObject *self, PyObject *unused)
{
    if (original_time)
        Py_RETURN_NONE;

    PyObject *datetime_module = PyImport_ImportModule("datetime");
    PyObject *datetime_class = PyObject_GetAttrString(datetime_module, "datetime");

    PyCFunctionObject *datetime_datetime_now = (PyCFunctionObject *) PyObject_GetAttrString(datetime_class, "now");
#if PY_VERSION_HEX >= 0x03070000
    original_now = (_PyCFunctionFastWithKeywords) datetime_datetime_now->m_ml->ml_meth;
#else
    original_now = (_PyCFunctionFast) datetime_datetime_now->m_ml->ml_meth;
#endif
    datetime_datetime_now->m_ml->ml_meth = (PyCFunction) _time_machine_now;
    Py_DECREF(datetime_datetime_now);

    PyCFunctionObject *datetime_datetime_utcnow = (PyCFunctionObject *) PyObject_GetAttrString(datetime_class, "utcnow");
    original_utcnow = datetime_datetime_utcnow->m_ml->ml_meth;
    datetime_datetime_utcnow->m_ml->ml_meth = _time_machine_utcnow;
    Py_DECREF(datetime_datetime_utcnow);

    Py_DECREF(datetime_class);
    Py_DECREF(datetime_module);

    PyObject *time_module = PyImport_ImportModule("time");

    PyCFunctionObject *time_time = (PyCFunctionObject *) PyObject_GetAttrString(time_module, "time");
    original_time = time_time->m_ml->ml_meth;
    time_time->m_ml->ml_meth = _time_machine_time;
    Py_DECREF(time_time);

    PyCFunctionObject *time_localtime = (PyCFunctionObject *) PyObject_GetAttrString(time_module, "localtime");
    original_localtime = time_localtime->m_ml->ml_meth;
    time_localtime->m_ml->ml_meth = _time_machine_localtime;
    Py_DECREF(time_localtime);

    PyCFunctionObject *time_gmtime = (PyCFunctionObject *) PyObject_GetAttrString(time_module, "gmtime");
    original_gmtime = time_gmtime->m_ml->ml_meth;
    time_gmtime->m_ml->ml_meth = _time_machine_gmtime;
    Py_DECREF(time_gmtime);

    PyCFunctionObject *time_strftime = (PyCFunctionObject *) PyObject_GetAttrString(time_module, "strftime");
    original_strftime = time_strftime->m_ml->ml_meth;
    time_strftime->m_ml->ml_meth = _time_machine_strftime;
    Py_DECREF(time_strftime);

    Py_DECREF(time_module);

    Py_RETURN_NONE;
}
PyDoc_STRVAR(patch_doc,
"patch() -> None\n\
\n\
Swap in helpers.");



PyDoc_STRVAR(module_doc, "_time_machine module");

static PyMethodDef module_methods[] = {
#if PY_VERSION_HEX >= 0x03070000
    {"original_now", (PyCFunction)_time_machine_original_now, METH_FASTCALL|METH_KEYWORDS, original_now_doc},
#else
    {"original_now", (PyCFunction)_time_machine_original_now, METH_FASTCALL, original_now_doc},
#endif
    {"original_utcnow", (PyCFunction)_time_machine_original_utcnow, METH_NOARGS, original_utcnow_doc},
    {"original_time", (PyCFunction)_time_machine_original_time, METH_NOARGS, original_time_doc},
    {"original_localtime", (PyCFunction)_time_machine_original_localtime, METH_VARARGS, original_localtime_doc},
    {"original_gmtime", (PyCFunction)_time_machine_original_gmtime, METH_VARARGS, original_gmtime_doc},
    {"original_strftime", (PyCFunction)_time_machine_original_strftime, METH_VARARGS, original_strftime_doc},
    {"patch", (PyCFunction)_time_machine_patch, METH_NOARGS, patch_doc},
    {NULL, NULL}  /* sentinel */
};

static struct PyModuleDef _time_machine_def = {
    PyModuleDef_HEAD_INIT,
    "_time_machine",
    module_doc,
    -1,
    module_methods,
    NULL,
    NULL,
    NULL,
    NULL
};

PyMODINIT_FUNC
PyInit__time_machine(void)
{
    PyObject *m;

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

    return m;
}

