#define PY_SSIZE_T_CLEAN
#define Py_CPYTHON_FRAMEOBJECT_H
#include "Python.h"
#include "marshal.h"
// #include "internal/pycore_frame.h"
#include "cpython/frameobject.h"

#define ARGTYPE 1
#include "libtomcrypt/headers/tomcrypt.h"

static const int MAJOR_VERSION = PY_MAJOR_VERSION;
static const int MINOR_VERSION = PY_MINOR_VERSION;

const unsigned char *aes_key = "WAIT_TO_BE_FILLED";

// PyObject *CodeObject;

PyObject* PyHumorRuntime_HumorEnter(PyObject* self, PyObject* args) {
    PyFrameObject *frame = PyEval_GetFrame();
    PyCodeObject *f_code = frame->f_code;
    // PyCodeObject *f_code = PyFrame_GetCode(PyEval_GetFrame());
    PyObject *codearray = f_code->co_code;

    Py_ssize_t size;
    unsigned char *data;
    PyBytes_AsStringAndSize(codearray, &data, &size);
    if (data == NULL) {
        return NULL;
    }
    if (size < 8 || size > 8192) {
        return NULL;
    }
    int use_size = (size - 8) - (size - 8) % 16;

    symmetric_key skey;
    int err;
    if ((err = aes_setup(aes_key + 1, 16, 0, &skey)) != CRYPT_OK) {
        return NULL;
    }
    for (int i = 0; i < use_size; i += 16) {
        if ((err = aes_ecb_decrypt(data + i + 8, data + i + 8, &skey)) != CRYPT_OK) {
            return NULL;
        }
    }
    aes_done(&skey);

    return Py_None;
}

PyObject* PyHumorRuntime_RunCode(PyObject* self, PyObject* args) {
    PyObject* name;
    PyObject* file;
    char* bytecode;
    Py_ssize_t bytecodeSize;
    if (!PyArg_ParseTuple(args, "OOy#", &name, &file, &bytecode, &bytecodeSize)) {
        return NULL;
    }

    // Parse the code meta
    if (bytecodeSize < 16) {
        return NULL;
    }
    if (bytecode[9] != MAJOR_VERSION || bytecode[10] != MINOR_VERSION) {
        return NULL;
    }
    if (*(int*)(bytecode + 12) != bytecodeSize - 16) {
        return NULL;
    }

    // Decrypt the bytecode
    symmetric_key skey;
    unsigned char *plaintext = malloc(bytecodeSize - 16);
    int err;
    if ((err = aes_setup(aes_key, 16, 0, &skey)) != CRYPT_OK) {
        return NULL;
    }
    for (int i = 0; i * 16 < bytecodeSize - 16; i++) {
        if ((err = aes_ecb_decrypt(bytecode + 16 + i * 16, plaintext + i * 16, &skey)) != CRYPT_OK) {
            return NULL;
        }
    }
    aes_done(&skey);

    // Load the code object
    PyObject *codeObject = PyMarshal_ReadObjectFromString(plaintext, bytecodeSize - 16);
    free(plaintext);
    if (codeObject == NULL) {
        return NULL;
    }

    // Add __humor_enter__ function to the globals
    PyMethodDef humorEnterMethodDef =
        {"__humor_enter__", PyHumorRuntime_HumorEnter, METH_VARARGS, NULL};

    PyObject* humorEnterFunc = PyCFunction_New(&humorEnterMethodDef, NULL);
    if (humorEnterFunc == NULL) {
        return NULL;
    }

    PyObject* globals = PyEval_GetGlobals();
    if (globals == NULL || PyDict_SetItemString(globals, "__humor_enter__", humorEnterFunc) < 0) {
        Py_DECREF(humorEnterFunc);
        return NULL;
    }

    Py_DECREF(humorEnterFunc);  // Decrease reference count since it's now in the globals

    // Run the code
    PyObject* result = PyImport_ExecCodeModuleObject(name, codeObject, file, NULL);
    Py_DECREF(codeObject);

    // Remove __humor_enter__ function from the globals
    PyDict_DelItemString(globals, "__humor_enter__");

    return result;
}

static PyMethodDef PyHumorRuntime_Methods[] = {
    {"__pyhumor__", PyHumorRuntime_RunCode, METH_VARARGS, NULL},
    {NULL, NULL, 0, NULL}
};

static struct PyModuleDef PyHumorRuntime_Module = {
    PyModuleDef_HEAD_INIT,
    "pyhumor_runtime",
    NULL,
    -1,
    PyHumorRuntime_Methods
};

PyMODINIT_FUNC PyInit_pyhumor_runtime(void) {
    return PyModule_Create(&PyHumorRuntime_Module);
}