/*
 * wraplib -- a complete wrapper around the fictional "tinylib" C library.
 *
 * Demonstrates the patterns every wrapper needs:
 *   - a heap type holding an opaque handle, with GC support
 *   - explicit close(), idempotent, plus __enter__/__exit__ and a
 *     tp_finalize backstop that emits a ResourceWarning
 *   - a handle validity check in every method (no use-after-close crashes)
 *   - one error-translation function mapping library codes to exceptions
 *   - a stored Python callback: strong reference, traversed, invoked with
 *     the thread state attached, exceptions routed to the unraisable hook
 *   - a PyCapsule exposing the raw handle without exposing a bare pointer
 *
 * Build:  pip install -e . --no-build-isolation
 * Lessons: 03, 04, 05, 06, 08, 10, 13
 */

#include <Python.h>
#include <limits.h>
#include <stddef.h>
#include <string.h>

#include "tinylib.h"

#define CAPSULE_NAME "wraplib.tiny_handle"

/* ================================================================== *
 * module state
 * ================================================================== */

typedef struct {
    PyObject *ConnType;
    PyObject *Error;        /* wraplib.Error, base of the module's errors */
} wraplib_state;

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

static PyModuleDef wraplib_module;   /* completed below; needed by slots */

/* One place that turns a library return code into a Python exception.
 * Returns 0 on success, -1 with an exception set otherwise. */
static int
check_rc(wraplib_state *state, int rc)
{
    if (rc == TINY_OK) {
        return 0;
    }
    switch (rc) {
    case TINY_ENOMEM:
        PyErr_NoMemory();
        break;
    case TINY_EINVAL:
        PyErr_SetString(PyExc_ValueError, tiny_strerror(rc));
        break;
    case TINY_ERANGE:
        PyErr_SetString(PyExc_OverflowError, tiny_strerror(rc));
        break;
    default:
        PyErr_Format(state->Error, "tinylib error %d: %s", rc, tiny_strerror(rc));
    }
    return -1;
}

/* ================================================================== *
 * the Conn object
 * ================================================================== */

typedef struct {
    PyObject_HEAD
    tiny_handle *h;         /* NULL once closed */
    PyObject    *name;      /* str, for repr and warnings */
    PyObject    *callback;  /* Python callable or NULL */
} ConnObject;

/* Every method starts here. Returning NULL means "already closed". */
static tiny_handle *
get_handle(ConnObject *self)
{
    if (self->h == NULL) {
        PyErr_SetString(PyExc_ValueError, "operation on a closed connection");
        return NULL;
    }
    return self->h;
}

static wraplib_state *
state_of(PyObject *self)
{
    PyObject *module = PyType_GetModuleByDef(Py_TYPE(self), &wraplib_module);
    if (module == NULL) {
        return NULL;
    }
    return get_state(module);   /* module is a borrowed reference */
}

/* ---- the trampoline the C library calls --------------------------- */

static void
conn_trampoline(void *userdata, int value)
{
    ConnObject *self = (ConnObject *)userdata;

    /* Attaching during finalization can hang or crash; never do it. */
    if (Py_IsFinalizing()) {
        return;
    }

    PyGILState_STATE gstate = PyGILState_Ensure();

    /* Take our own strong reference: another thread could replace the
     * callback attribute while we are calling it. */
    PyObject *cb = Py_XNewRef(self->callback);
    if (cb != NULL) {
        PyObject *arg = PyLong_FromLong(value);
        PyObject *res = (arg != NULL) ? PyObject_CallOneArg(cb, arg) : NULL;
        if (res == NULL) {
            /* No caller to propagate to: report rather than swallow. */
            PyErr_FormatUnraisable("Exception ignored in %R callback",
                                   (PyObject *)self);
        }
        Py_XDECREF(res);
        Py_XDECREF(arg);
        Py_DECREF(cb);
    }

    PyGILState_Release(gstate);
}

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

static PyObject *
Conn_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
    ConnObject *self = (ConnObject *)type->tp_alloc(type, 0);
    if (self == NULL) {
        return NULL;
    }
    self->h = NULL;
    self->name = NULL;
    self->callback = NULL;
    return (PyObject *)self;
}

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

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "U:Conn", kwlist, &name)) {
        return -1;
    }

    wraplib_state *state = state_of(op);
    if (state == NULL) {
        return -1;
    }

    /* PyUnicode_AsUTF8AndSize returns a pointer owned by `name`; it is
     * valid here because `name` is alive and we stay attached. */
    Py_ssize_t len;
    const char *cname = PyUnicode_AsUTF8AndSize(name, &len);
    if (cname == NULL) {
        return -1;
    }
    if ((Py_ssize_t)strlen(cname) != len) {
        PyErr_SetString(PyExc_ValueError, "name must not contain NUL");
        return -1;
    }

    tiny_handle *h = NULL;
    if (check_rc(state, tiny_open(cname, &h)) < 0) {
        return -1;
    }

    /* Re-initialization on an already-open object: close the old handle. */
    if (self->h != NULL) {
        tiny_close(self->h);
    }
    self->h = h;
    Py_XSETREF(self->name, Py_NewRef(name));

    /* We hand `self` to the library as raw userdata, which does not keep the
     * object alive. That is safe here only because the handle is owned by
     * this object and is closed (clearing the callback) before the object
     * dies. A library that could outlive the registration would require an
     * explicit Py_INCREF held for the duration. */
    tiny_set_callback(self->h, conn_trampoline, self);
    return 0;
}

/* The single close path: idempotent and safe to call from a finalizer. */
static void
conn_do_close(ConnObject *self)
{
    tiny_handle *h = self->h;
    if (h == NULL) {
        return;
    }
    self->h = NULL;          /* NULL first: makes close() re-entrant-safe */
    tiny_set_callback(h, NULL, NULL);

    Py_BEGIN_ALLOW_THREADS   /* a real library's close may block */
    tiny_close(h);
    Py_END_ALLOW_THREADS
}

static int
Conn_traverse(PyObject *op, visitproc visit, void *arg)
{
    ConnObject *self = (ConnObject *)op;
    Py_VISIT(Py_TYPE(self));
    Py_VISIT(self->name);
    Py_VISIT(self->callback);   /* a closure here can reference self */
    return 0;
}

static int
Conn_clear(PyObject *op)
{
    ConnObject *self = (ConnObject *)op;
    Py_CLEAR(self->name);
    Py_CLEAR(self->callback);
    return 0;
}

/* tp_finalize runs once, while the object is still valid. Non-Python
 * resources and user-visible warnings belong here, not in tp_clear. */
static void
Conn_finalize(PyObject *op)
{
    ConnObject *self = (ConnObject *)op;
    if (self->h == NULL) {
        return;
    }
    /* A finalizer can run with an exception in flight. */
    PyObject *exc = PyErr_GetRaisedException();

    if (PyErr_ResourceWarning(op, 1, "unclosed connection %R", op) < 0) {
        PyErr_Clear();   /* warnings-as-errors here is not actionable */
    }
    conn_do_close(self);

    PyErr_SetRaisedException(exc);
}

static void
Conn_dealloc(PyObject *op)
{
    PyTypeObject *tp = Py_TYPE(op);

    PyObject_GC_UnTrack(op);
    if (PyObject_CallFinalizerFromDealloc(op) < 0) {
        return;   /* the object was resurrected by its finalizer */
    }
    conn_do_close((ConnObject *)op);   /* belt and braces */
    (void)Conn_clear(op);
    tp->tp_free(op);
    Py_DECREF(tp);
}

/* ---- methods ------------------------------------------------------ */

static PyObject *
Conn_push(PyObject *op, PyObject *arg)
{
    ConnObject *self = (ConnObject *)op;
    tiny_handle *h = get_handle(self);
    if (h == NULL) {
        return NULL;
    }
    long v = PyLong_AsLong(arg);
    if (v == -1 && PyErr_Occurred()) {
        return NULL;
    }
    if (v < INT_MIN || v > INT_MAX) {
        PyErr_SetString(PyExc_OverflowError, "value does not fit in a C int");
        return NULL;
    }

    wraplib_state *state = state_of(op);
    if (state == NULL) {
        return NULL;
    }
    /* NOTE: tiny_push invokes our callback synchronously, which runs Python
     * code, so we must stay attached here. A library call that does NOT call
     * back would be wrapped in Py_BEGIN_ALLOW_THREADS instead. */
    if (check_rc(state, tiny_push(h, (int)v)) < 0) {
        return NULL;
    }
    Py_RETURN_NONE;
}

static PyObject *
Conn_close(PyObject *op, PyObject *Py_UNUSED(ignored))
{
    conn_do_close((ConnObject *)op);
    Py_RETURN_NONE;
}

static PyObject *
Conn_enter(PyObject *op, PyObject *Py_UNUSED(ignored))
{
    if (get_handle((ConnObject *)op) == NULL) {
        return NULL;
    }
    return Py_NewRef(op);
}

static PyObject *
Conn_exit(PyObject *op, PyObject *const *args, Py_ssize_t nargs)
{
    conn_do_close((ConnObject *)op);
    Py_RETURN_FALSE;    /* do not suppress exceptions */
}

/* Hand the raw handle to other C code without exposing a bare pointer.
 * The capsule does NOT own the handle; the Conn object still does. */
static PyObject *
Conn_capsule(PyObject *op, PyObject *Py_UNUSED(ignored))
{
    tiny_handle *h = get_handle((ConnObject *)op);
    if (h == NULL) {
        return NULL;
    }
    return PyCapsule_New(h, CAPSULE_NAME, NULL);
}

static PyMethodDef Conn_methods[] = {
    {"push", Conn_push, METH_O,
     PyDoc_STR("push(value, /)\n--\n\nAdd a value; invokes the callback.")},
    {"close", Conn_close, METH_NOARGS,
     PyDoc_STR("close()\n--\n\nRelease the handle. Idempotent.")},
    {"__enter__", Conn_enter, METH_NOARGS, NULL},
    {"__exit__", (PyCFunction)(void (*)(void))Conn_exit, METH_FASTCALL, NULL},
    {"as_capsule", Conn_capsule, METH_NOARGS,
     PyDoc_STR("as_capsule()\n--\n\nReturn the raw handle wrapped in a capsule.")},
    {NULL}
};

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

static PyObject *
Conn_get_total(PyObject *op, void *closure)
{
    tiny_handle *h = get_handle((ConnObject *)op);
    if (h == NULL) {
        return NULL;
    }
    wraplib_state *state = state_of(op);
    if (state == NULL) {
        return NULL;
    }
    long total = 0;
    if (check_rc(state, tiny_total(h, &total)) < 0) {
        return NULL;
    }
    return PyLong_FromLong(total);
}

static PyObject *
Conn_get_closed(PyObject *op, void *closure)
{
    return PyBool_FromLong(((ConnObject *)op)->h == NULL);
}

static PyObject *
Conn_get_callback(PyObject *op, void *closure)
{
    ConnObject *self = (ConnObject *)op;
    return self->callback ? Py_NewRef(self->callback) : Py_NewRef(Py_None);
}

static int
Conn_set_callback(PyObject *op, PyObject *value, void *closure)
{
    ConnObject *self = (ConnObject *)op;
    if (value == NULL || value == Py_None) {
        Py_CLEAR(self->callback);
        return 0;
    }
    if (!PyCallable_Check(value)) {
        PyErr_SetString(PyExc_TypeError, "callback must be callable or None");
        return -1;
    }
    Py_XSETREF(self->callback, Py_NewRef(value));
    return 0;
}

static PyGetSetDef Conn_getset[] = {
    {"total", Conn_get_total, NULL, PyDoc_STR("running total"), NULL},
    {"closed", Conn_get_closed, NULL, PyDoc_STR("True once closed"), NULL},
    {"callback", Conn_get_callback, Conn_set_callback,
     PyDoc_STR("callable invoked on each push, or None"), NULL},
    {NULL}
};

static PyObject *
Conn_repr(PyObject *op)
{
    ConnObject *self = (ConnObject *)op;
    return PyUnicode_FromFormat("<%s %R%s>", Py_TYPE(self)->tp_name,
                                self->name ? self->name : Py_None,
                                self->h ? "" : " (closed)");
}

static PyType_Slot Conn_slots[] = {
    {Py_tp_doc,      (void *)PyDoc_STR("Conn(name)\n--\n\nA tinylib connection.")},
    {Py_tp_new,      Conn_new},
    {Py_tp_init,     Conn_init},
    {Py_tp_dealloc,  Conn_dealloc},
    {Py_tp_finalize, Conn_finalize},
    {Py_tp_traverse, Conn_traverse},
    {Py_tp_clear,    Conn_clear},
    {Py_tp_repr,     Conn_repr},
    {Py_tp_methods,  Conn_methods},
    {Py_tp_getset,   Conn_getset},
    {0, NULL}
};

static PyType_Spec Conn_spec = {
    .name = "wraplib.Conn",
    .basicsize = sizeof(ConnObject),
    .itemsize = 0,
    .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,   /* no BASETYPE: the
                 close/finalize invariants are not designed for subclassing */
    .slots = Conn_slots,
};

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

static int
wraplib_exec(PyObject *module)
{
    wraplib_state *st = get_state(module);

    st->Error = PyErr_NewExceptionWithDoc("wraplib.Error",
                                          "Base error for wraplib.", NULL, NULL);
    if (st->Error == NULL) {
        return -1;
    }
    if (PyModule_AddObjectRef(module, "Error", st->Error) < 0) {
        return -1;
    }

    st->ConnType = PyType_FromModuleAndSpec(module, &Conn_spec, NULL);
    if (st->ConnType == NULL) {
        return -1;
    }
    if (PyModule_AddObjectRef(module, "Conn", st->ConnType) < 0) {
        return -1;
    }

    /* Export the library's constants under their own C names. */
    if (PyModule_AddIntMacro(module, TINY_OK) < 0) return -1;
    if (PyModule_AddIntMacro(module, TINY_EINVAL) < 0) return -1;
    if (PyModule_AddIntMacro(module, TINY_ENOMEM) < 0) return -1;
    if (PyModule_AddIntMacro(module, TINY_ERANGE) < 0) return -1;
    return 0;
}

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

static int
wraplib_clear(PyObject *module)
{
    wraplib_state *st = get_state(module);
    Py_CLEAR(st->ConnType);
    Py_CLEAR(st->Error);
    return 0;
}

static void
wraplib_free(void *module)
{
    (void)wraplib_clear((PyObject *)module);
}

static PyModuleDef_Slot wraplib_slots[] = {
    {Py_mod_exec, wraplib_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 wraplib_module = {
    .m_base     = PyModuleDef_HEAD_INIT,
    .m_name     = "wraplib",
    .m_doc      = PyDoc_STR("Wrapper around the fictional tinylib."),
    .m_size     = sizeof(wraplib_state),
    .m_slots    = wraplib_slots,
    .m_traverse = wraplib_traverse,
    .m_clear    = wraplib_clear,
    .m_free     = wraplib_free,
};

PyMODINIT_FUNC
PyInit_wraplib(void)
{
    return PyModuleDef_Init(&wraplib_module);
}