#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <structmember.h>

// forward-declare our function
static PyObject* fast_query_many(PyObject *self, PyObject *args, PyObject *kwargs);

// Method table
typedef struct {
    const char *name;
    PyCFunction func;
    int flags;
    const char *doc;
} MethodDef;

static PyMethodDef FastQueryMethods[] = {
    {"query_many", (PyCFunction)fast_query_many, METH_VARARGS | METH_KEYWORDS,
     "query_many(hashes: list[tuple[int,int]], index: dict, threshold: int=1) -> list[tuple[(song_id, delta), count]] (song_id may be str or int)"},
    {NULL, NULL, 0, NULL}
};

// Module definition
static struct PyModuleDef fastquerymodule = {
    PyModuleDef_HEAD_INIT,
    "fast_query",   /* name of module */
    "High-speed query_many implemented in C with string/int support for song_id", /* module doc */
    -1,             /* size of per-interpreter state or -1 */
    FastQueryMethods
};

// Module init
PyMODINIT_FUNC PyInit_fast_query(void) {
    return PyModule_Create(&fastquerymodule);
}

static PyObject* fast_query_many(PyObject *self, PyObject *args, PyObject *kwargs) {
    PyObject *hashes = NULL, *index = NULL;
    int threshold = 1;
    static char *kwlist[] = {"hashes","index","threshold", NULL};
    
    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OO|i", kwlist,
                                     &hashes, &index, &threshold))
        return NULL;

    // Convert hashes to sequence
    PyObject *seq_hashes = PySequence_Fast(hashes, "hashes must be a sequence");
    if (!seq_hashes) return NULL;
    Py_ssize_t n_hashes = PySequence_Fast_GET_SIZE(seq_hashes);

    // counts dict
    PyObject *counts = PyDict_New();
    if (!counts) { Py_DECREF(seq_hashes); return NULL; }

    for (Py_ssize_t i = 0; i < n_hashes; i++) {
        PyObject *pair = PySequence_Fast_GET_ITEM(seq_hashes, i);
        long h = PyLong_AsLong(PyTuple_GET_ITEM(pair, 0));
        long t0 = PyLong_AsLong(PyTuple_GET_ITEM(pair, 1));

        // lookup bucket = index[h]
        PyObject *key_h = PyLong_FromLong(h);
        if (!key_h) goto error;
        PyObject *bucket = PyObject_GetItem(index, key_h);
        Py_DECREF(key_h);
        if (!bucket) { PyErr_Clear(); continue; }
        if (!PyList_Check(bucket)) { Py_DECREF(bucket); continue; }
        Py_ssize_t bucket_sz = PyList_Size(bucket);

        for (Py_ssize_t j = 0; j < bucket_sz; j++) {
            PyObject *songpair = PyList_GetItem(bucket, j);  // borrowed
            PyObject *song_id_obj = PyTuple_GET_ITEM(songpair, 0);
            PyObject *off_list = PyTuple_GET_ITEM(songpair, 1);
            Py_ssize_t off_sz = PyList_Size(off_list);

            for (Py_ssize_t k = 0; k < off_sz; k++) {
                long off = PyLong_AsLong(PyList_GetItem(off_list, k));
                long delta = off - t0;

                // build key tuple
                PyObject *res_key = PyTuple_New(2);
                if (!res_key) goto error;
                Py_INCREF(song_id_obj);
                PyTuple_SET_ITEM(res_key, 0, song_id_obj);
                PyObject *delta_obj = PyLong_FromLong(delta);
                if (!delta_obj) { Py_DECREF(res_key); goto error; }
                PyTuple_SET_ITEM(res_key, 1, delta_obj);

                // increment counts
                PyObject *old = PyDict_GetItem(counts, res_key);
                PyObject *newcount;
                if (old) {
                    long c = PyLong_AsLong(old) + 1;
                    newcount = PyLong_FromLong(c);
                } else {
                    newcount = PyLong_FromLong(1);
                }
                if (!newcount) { Py_DECREF(res_key); goto error; }

                if (PyDict_SetItem(counts, res_key, newcount) < 0) {
                    Py_DECREF(res_key);
                    Py_DECREF(newcount);
                    goto error;
                }
                Py_DECREF(res_key);
                Py_DECREF(newcount);
            }
        }
        Py_DECREF(bucket);
    }
    Py_DECREF(seq_hashes);

    // delegate top-k to collections.Counter
    PyObject *collections = PyImport_ImportModule("collections");
    if (!collections) goto error;
    PyObject *CounterCls = PyObject_GetAttrString(collections, "Counter");
    Py_DECREF(collections);
    if (!CounterCls) goto error;

    PyObject *counter = PyObject_CallFunctionObjArgs(CounterCls, counts, NULL);
    Py_DECREF(CounterCls);
    if (!counter) goto error;
    Py_DECREF(counts);

    PyObject *most_common = PyObject_GetAttrString(counter, "most_common");
    if (!most_common) { Py_DECREF(counter); goto error; }
    PyObject *result = PyObject_CallFunction(most_common, "i", 100);
    Py_DECREF(most_common);
    Py_DECREF(counter);
    return result;

error:
    Py_XDECREF(counts);
    Py_XDECREF(seq_hashes);
    return NULL;
}