/*
 * counter -- a heap type with per-module state and full cycle-GC support.
 *
 * Demonstrates:
 *   - per-module state instead of C globals, with m_traverse/m_clear/m_free
 *   - a heap type created with PyType_Spec / PyType_FromModuleAndSpec
 *   - the full lifecycle: tp_new, tp_init, tp_traverse, tp_clear, tp_dealloc
 *   - members, getset with validation, and __repr__
 *   - reaching module state from a method via METH_METHOD and defining_class
 *
 * Build:  pip install -e . --no-build-isolation
 * Lessons: 03, 04, 05, 06, 08
 */

#include <Python.h>
#include <stddef.h>   /* offsetof */

/* ================================================================== *
 * per-module state
 * ================================================================== */

typedef struct {
    PyObject *CounterType;   /* strong ref to our heap type   */
    PyObject *Error;         /* strong ref to counter.Error   */
} counter_state;

static inline counter_state *
get_state(PyObject *module)
{
    void *state = PyModule_GetState(module);
    assert(state != NULL);
    return (counter_state *)state;
}

/* Needed by slot functions, which receive neither the module nor the
 * defining class. Declared here, defined after the module definition. */
static PyModuleDef counter_module;

/* ================================================================== *
 * the instance
 * ================================================================== */

typedef struct {
    PyObject_HEAD
    PyObject *name;       /* str, never NULL after tp_new succeeds */
    PyObject *payload;    /* arbitrary object or NULL              */
    long      count;      /* plain C, never traversed              */
} CounterObject;

/* ---- lifecycle --------------------------------------------------- */

static PyObject *
Counter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
    /* Allocate through the *passed* type so subclasses work. */
    CounterObject *self = (CounterObject *)type->tp_alloc(type, 0);
    if (self == NULL) {
        return NULL;
    }
    /* The object is GC-tracked from here on, so every PyObject* field must
     * be in a state tp_traverse can cope with -- NULL is fine, garbage is not. */
    self->payload = NULL;
    self->count = 0;
    self->name = PyUnicode_FromString("");
    if (self->name == NULL) {
        Py_DECREF(self);
        return NULL;
    }
    return (PyObject *)self;
}

static int
Counter_init(PyObject *op, PyObject *args, PyObject *kwds)
{
    CounterObject *self = (CounterObject *)op;
    static char *kwlist[] = {"name", "payload", NULL};
    PyObject *name = NULL;      /* borrowed */
    PyObject *payload = NULL;   /* borrowed */

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|UO:Counter", kwlist,
                                     &name, &payload)) {
        return -1;
    }
    if (name != NULL) {
        Py_SETREF(self->name, Py_NewRef(name));
    }
    if (payload != NULL) {
        Py_XSETREF(self->payload, Py_NewRef(payload));
    }
    self->count = 0;
    return 0;
}

static int
Counter_traverse(PyObject *op, visitproc visit, void *arg)
{
    CounterObject *self = (CounterObject *)op;
    /* Heap types: an instance references its type, so the type must be
     * visited or the module can never be collected. */
    Py_VISIT(Py_TYPE(self));
    Py_VISIT(self->name);
    Py_VISIT(self->payload);
    return 0;
}

static int
Counter_clear(PyObject *op)
{
    CounterObject *self = (CounterObject *)op;
    Py_CLEAR(self->name);
    Py_CLEAR(self->payload);
    return 0;
}

static void
Counter_dealloc(PyObject *op)
{
    PyTypeObject *tp = Py_TYPE(op);   /* read before the memory goes away */
    PyObject_GC_UnTrack(op);          /* 1. the collector must stop seeing us */
    (void)Counter_clear(op);          /* 2. drop object references            */
    tp->tp_free(op);                  /* 3. release the memory                */
    Py_DECREF(tp);                    /* 4. release this instance's ref to the type */
}

/* ---- behaviour --------------------------------------------------- */

static PyObject *
Counter_repr(PyObject *op)
{
    CounterObject *self = (CounterObject *)op;
    /* Fields can be NULL after tp_clear, so guard. */
    if (self->name == NULL) {
        return PyUnicode_FromFormat("<%s (cleared)>", Py_TYPE(self)->tp_name);
    }
    return PyUnicode_FromFormat("<%s %R count=%ld>",
                                Py_TYPE(self)->tp_name, self->name, self->count);
}

/* bump(n=1) -> int
 *
 * METH_METHOD | METH_FASTCALL | METH_KEYWORDS gives us `defining_class`,
 * which is how a heap-type method reaches per-module state. */
static PyObject *
Counter_bump(PyObject *op, PyTypeObject *defining_class,
             PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
{
    CounterObject *self = (CounterObject *)op;
    long n = 1;

    if (nargs > 1) {
        PyErr_SetString(PyExc_TypeError,
                        "bump() takes at most 1 positional argument");
        return NULL;
    }
    if (nargs == 1) {
        n = PyLong_AsLong(args[0]);
        if (n == -1 && PyErr_Occurred()) {
            return NULL;
        }
    }
    if (kwnames != NULL) {
        Py_ssize_t nkw = PyTuple_GET_SIZE(kwnames);
        for (Py_ssize_t i = 0; i < nkw; i++) {
            PyObject *key = PyTuple_GET_ITEM(kwnames, i);   /* tuple: safe */
            if (PyUnicode_CompareWithASCIIString(key, "n") != 0) {
                PyErr_Format(PyExc_TypeError,
                             "bump() got an unexpected keyword argument %R", key);
                return NULL;
            }
            if (nargs == 1) {
                PyErr_SetString(PyExc_TypeError,
                                "bump() got multiple values for argument 'n'");
                return NULL;
            }
            n = PyLong_AsLong(args[nargs + i]);
            if (n == -1 && PyErr_Occurred()) {
                return NULL;
            }
        }
    }

    if (n < 0) {
        counter_state *st = (counter_state *)PyType_GetModuleState(defining_class);
        if (st == NULL) {
            return NULL;
        }
        PyErr_SetString(st->Error, "bump() amount must be non-negative");
        return NULL;
    }

    self->count += n;
    return PyLong_FromLong(self->count);
}

static PyObject *
Counter_reset(PyObject *op, PyObject *Py_UNUSED(ignored))
{
    ((CounterObject *)op)->count = 0;
    Py_RETURN_NONE;
}

/* A slot function gets neither module nor defining class, so it goes
 * through the module definition to find its module. */
static PyObject *
Counter_richcompare(PyObject *a, PyObject *b, int op)
{
    if (op != Py_EQ && op != Py_NE) {
        Py_RETURN_NOTIMPLEMENTED;
    }
    PyObject *module = PyType_GetModuleByDef(Py_TYPE(a), &counter_module);
    if (module == NULL) {
        return NULL;   /* not one of ours */
    }
    counter_state *st = get_state(module);   /* module is borrowed */
    if (!PyObject_TypeCheck(b, (PyTypeObject *)st->CounterType)) {
        Py_RETURN_NOTIMPLEMENTED;
    }
    int equal = ((CounterObject *)a)->count == ((CounterObject *)b)->count;
    if (op == Py_NE) {
        equal = !equal;
    }
    return PyBool_FromLong(equal);
}

/* ---- attributes -------------------------------------------------- */

static PyMemberDef Counter_members[] = {
    {"count", Py_T_LONG, offsetof(CounterObject, count), Py_READONLY,
     PyDoc_STR("current count")},
    {NULL}
};

static PyObject *
Counter_get_name(PyObject *op, void *closure)
{
    CounterObject *self = (CounterObject *)op;
    if (self->name == NULL) {
        PyErr_SetString(PyExc_AttributeError, "name");
        return NULL;
    }
    return Py_NewRef(self->name);
}

static int
Counter_set_name(PyObject *op, PyObject *value, void *closure)
{
    CounterObject *self = (CounterObject *)op;
    if (value == NULL) {
        PyErr_SetString(PyExc_AttributeError, "cannot delete 'name'");
        return -1;
    }
    if (!PyUnicode_Check(value)) {
        PyErr_Format(PyExc_TypeError, "name must be str, not %s",
                     Py_TYPE(value)->tp_name);
        return -1;
    }
    Py_XSETREF(self->name, Py_NewRef(value));
    return 0;
}

static PyObject *
Counter_get_payload(PyObject *op, void *closure)
{
    CounterObject *self = (CounterObject *)op;
    if (self->payload == NULL) {
        Py_RETURN_NONE;
    }
    return Py_NewRef(self->payload);
}

static int
Counter_set_payload(PyObject *op, PyObject *value, void *closure)
{
    CounterObject *self = (CounterObject *)op;
    Py_XSETREF(self->payload, Py_XNewRef(value));   /* NULL means delete */
    return 0;
}

static PyGetSetDef Counter_getset[] = {
    {"name", Counter_get_name, Counter_set_name,
     PyDoc_STR("the counter's name (str)"), NULL},
    {"payload", Counter_get_payload, Counter_set_payload,
     PyDoc_STR("arbitrary attached object"), NULL},
    {NULL}
};

static PyMethodDef Counter_methods[] = {
    {"bump", (PyCFunction)(void (*)(void))Counter_bump,
     METH_METHOD | METH_FASTCALL | METH_KEYWORDS,
     PyDoc_STR("bump(n=1)\n--\n\nAdd n to the count and return the new value.")},
    {"reset", Counter_reset, METH_NOARGS,
     PyDoc_STR("reset()\n--\n\nSet the count back to zero.")},
    {NULL}
};

static PyType_Slot Counter_slots[] = {
    {Py_tp_doc,         (void *)PyDoc_STR("Counter(name='', payload=None)")},
    {Py_tp_new,         Counter_new},
    {Py_tp_init,        Counter_init},
    {Py_tp_dealloc,     Counter_dealloc},
    {Py_tp_traverse,    Counter_traverse},
    {Py_tp_clear,       Counter_clear},
    {Py_tp_repr,        Counter_repr},
    {Py_tp_richcompare, Counter_richcompare},
    {Py_tp_methods,     Counter_methods},
    {Py_tp_members,     Counter_members},
    {Py_tp_getset,      Counter_getset},
    {0, NULL}
};

static PyType_Spec Counter_spec = {
    .name = "counter.Counter",
    .basicsize = sizeof(CounterObject),
    .itemsize = 0,
    .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
    .slots = Counter_slots,
};

/* ================================================================== *
 * module-level function: `self` is the module
 * ================================================================== */

static PyObject *
counter_make(PyObject *module, PyObject *arg)
{
    counter_state *st = get_state(module);
    if (!PyUnicode_Check(arg)) {
        PyErr_SetString(st->Error, "make() requires a str name");
        return NULL;
    }
    return PyObject_CallOneArg(st->CounterType, arg);
}

static PyMethodDef counter_methods[] = {
    {"make", counter_make, METH_O,
     PyDoc_STR("make(name, /)\n--\n\nCreate a Counter with the given name.")},
    {NULL}
};

/* ================================================================== *
 * module plumbing
 * ================================================================== */

static int
counter_exec(PyObject *module)
{
    counter_state *st = get_state(module);   /* already zeroed */

    st->Error = PyErr_NewExceptionWithDoc("counter.Error",
                                          "Errors raised by the counter module.",
                                          NULL, NULL);
    if (st->Error == NULL) {
        return -1;
    }
    if (PyModule_AddObjectRef(module, "Error", st->Error) < 0) {
        return -1;
    }

    st->CounterType = PyType_FromModuleAndSpec(module, &Counter_spec, NULL);
    if (st->CounterType == NULL) {
        return -1;
    }
    if (PyModule_AddObjectRef(module, "Counter", st->CounterType) < 0) {
        return -1;
    }
    return 0;
}

static int
counter_traverse(PyObject *module, visitproc visit, void *arg)
{
    counter_state *st = get_state(module);
    Py_VISIT(st->CounterType);
    Py_VISIT(st->Error);
    return 0;
}

static int
counter_clear(PyObject *module)
{
    counter_state *st = get_state(module);
    Py_CLEAR(st->CounterType);
    Py_CLEAR(st->Error);
    return 0;
}

static void
counter_free(void *module)
{
    (void)counter_clear((PyObject *)module);
}

static PyModuleDef_Slot counter_slots[] = {
    {Py_mod_exec, counter_exec},
    {Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED},
#if PY_VERSION_HEX >= 0x030D0000
    {Py_mod_gil, Py_MOD_GIL_NOT_USED},
#endif
    {0, NULL}
};

static PyModuleDef counter_module = {
    .m_base     = PyModuleDef_HEAD_INIT,
    .m_name     = "counter",
    .m_doc      = PyDoc_STR("Heap type with per-module state."),
    .m_size     = sizeof(counter_state),
    .m_methods  = counter_methods,
    .m_slots    = counter_slots,
    .m_traverse = counter_traverse,
    .m_clear    = counter_clear,
    .m_free     = counter_free,
};

PyMODINIT_FUNC
PyInit_counter(void)
{
    return PyModuleDef_Init(&counter_module);
}