#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <stdint.h>
#include "CPy.h"
#define NATIVE_INTERNAL_MODULE
#include "native_internal.h"

#define START_SIZE 512
#define MAX_SHORT_INT_TAGGED (255 << 1)

#define MAX_SHORT_LEN 127
#define LONG_STR_TAG 1

#define MIN_SHORT_INT -10
#define MAX_SHORT_INT 117
#define MEDIUM_INT_TAG 1
#define LONG_INT_TAG 3

#define CPY_BOOL_ERROR 2
#define CPY_NONE_ERROR 2
#define CPY_NONE 1

#define _CHECK_BUFFER(data, err)      if (unlikely(_check_buffer(data) == CPY_NONE_ERROR)) \
                                          return err;
#define _CHECK_SIZE(data, need)       if (unlikely(_check_size((BufferObject *)data, need) == CPY_NONE_ERROR)) \
                                          return CPY_NONE_ERROR;
#define _CHECK_READ(data, size, err)  if (unlikely(_check_read((BufferObject *)data, size) == CPY_NONE_ERROR)) \
                                          return err;

#define _READ(data, type)  *(type *)(((BufferObject *)data)->buf + ((BufferObject *)data)->pos); \
                           ((BufferObject *)data)->pos += sizeof(type);

#define _WRITE(data, type, v)  *(type *)(((BufferObject *)data)->buf + ((BufferObject *)data)->pos) = v; \
                               ((BufferObject *)data)->pos += sizeof(type);

typedef struct {
    PyObject_HEAD
    Py_ssize_t pos;
    Py_ssize_t end;
    Py_ssize_t size;
    char *buf;
    PyObject *source;
} BufferObject;

static PyTypeObject BufferType;

static PyObject*
Buffer_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
    if (type != &BufferType) {
        PyErr_SetString(PyExc_TypeError, "Buffer should not be subclassed");
        return NULL;
    }

    BufferObject *self = (BufferObject *)type->tp_alloc(type, 0);
    if (self != NULL) {
        self->pos = 0;
        self->end = 0;
        self->size = 0;
        self->buf = NULL;
    }
    return (PyObject *) self;
}


static int
Buffer_init_internal(BufferObject *self, PyObject *source) {
    if (source) {
        if (!PyBytes_Check(source)) {
            PyErr_SetString(PyExc_TypeError, "source must be a bytes object");
            return -1;
        }
        self->size = PyBytes_GET_SIZE(source);
        self->end = self->size;
        // This returns a pointer to internal bytes data, so make our own copy.
        char *buf = PyBytes_AsString(source);
        self->buf = PyMem_Malloc(self->size);
        memcpy(self->buf, buf, self->size);
    } else {
        self->buf = PyMem_Malloc(START_SIZE);
        self->size = START_SIZE;
    }
    return 0;
}

static PyObject*
Buffer_internal(PyObject *source) {
    BufferObject *self = (BufferObject *)BufferType.tp_alloc(&BufferType, 0);
    if (self == NULL)
        return NULL;
    self->pos = 0;
    self->end = 0;
    self->size = 0;
    self->buf = NULL;
    if (Buffer_init_internal(self, source) == -1) {
        Py_DECREF(self);
        return NULL;
    }
    return (PyObject *)self;
}

static PyObject*
Buffer_internal_empty(void) {
    return Buffer_internal(NULL);
}

static int
Buffer_init(BufferObject *self, PyObject *args, PyObject *kwds)
{
    static char *kwlist[] = {"source", NULL};
    PyObject *source = NULL;

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O", kwlist, &source))
        return -1;

    return Buffer_init_internal(self, source);
}

static void
Buffer_dealloc(BufferObject *self)
{
    PyMem_Free(self->buf);
    Py_TYPE(self)->tp_free((PyObject *)self);
}

static PyObject*
Buffer_getvalue_internal(PyObject *self)
{
    return PyBytes_FromStringAndSize(((BufferObject *)self)->buf, ((BufferObject *)self)->end);
}

static PyObject*
Buffer_getvalue(BufferObject *self, PyObject *Py_UNUSED(ignored))
{
    return PyBytes_FromStringAndSize(self->buf, self->end);
}

static PyMethodDef Buffer_methods[] = {
    {"getvalue", (PyCFunction) Buffer_getvalue, METH_NOARGS,
     "Return the buffer content as bytes object"
    },
    {NULL}  /* Sentinel */
};

static PyTypeObject BufferType = {
    .ob_base = PyVarObject_HEAD_INIT(NULL, 0)
    .tp_name = "Buffer",
    .tp_doc = PyDoc_STR("Mypy cache buffer objects"),
    .tp_basicsize = sizeof(BufferObject),
    .tp_itemsize = 0,
    .tp_flags = Py_TPFLAGS_DEFAULT,
    .tp_new = Buffer_new,
    .tp_init = (initproc) Buffer_init,
    .tp_dealloc = (destructor) Buffer_dealloc,
    .tp_methods = Buffer_methods,
};

static inline char
_check_buffer(PyObject *data) {
    if (unlikely(Py_TYPE(data) != &BufferType)) {
        PyErr_Format(
            PyExc_TypeError, "data must be a Buffer object, got %s", Py_TYPE(data)->tp_name
        );
        return CPY_NONE_ERROR;
    }
    return CPY_NONE;
}

static inline char
_check_size(BufferObject *data, Py_ssize_t need) {
    Py_ssize_t target = data->pos + need;
    if (target <= data->size)
        return CPY_NONE;
    do
        data->size *= 2;
    while (target >= data->size);
    data->buf = PyMem_Realloc(data->buf, data->size);
    if (unlikely(data->buf == NULL)) {
        PyErr_NoMemory();
        return CPY_NONE_ERROR;
    }
    return CPY_NONE;
}

static inline char
_check_read(BufferObject *data, Py_ssize_t need) {
    if (unlikely(data->pos + need > data->end)) {
        PyErr_SetString(PyExc_ValueError, "reading past the buffer end");
        return CPY_NONE_ERROR;
    }
    return CPY_NONE;
}

/*
bool format: single byte
    \x00 - False
    \x01 - True
*/

static char
read_bool_internal(PyObject *data) {
    _CHECK_BUFFER(data, CPY_BOOL_ERROR)
    _CHECK_READ(data, 1, CPY_BOOL_ERROR)
    char res = _READ(data, char)
    return res;
}

static PyObject*
read_bool(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames) {
    static const char * const kwlist[] = {"data", 0};
    static CPyArg_Parser parser = {"O:read_bool", kwlist, 0};
    PyObject *data;
    if (unlikely(!CPyArg_ParseStackAndKeywordsOneArg(args, nargs, kwnames, &parser, &data))) {
        return NULL;
    }
    char res = read_bool_internal(data);
    if (unlikely(res == CPY_BOOL_ERROR))
        return NULL;
    PyObject *retval = res ? Py_True : Py_False;
    Py_INCREF(retval);
    return retval;
}

static char
write_bool_internal(PyObject *data, char value) {
    _CHECK_BUFFER(data, CPY_NONE_ERROR)
    _CHECK_SIZE(data, 1)
    _WRITE(data, char, value)
    ((BufferObject *)data)->end += 1;
    return CPY_NONE;
}

static PyObject*
write_bool(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames) {
    static const char * const kwlist[] = {"data", "value", 0};
    static CPyArg_Parser parser = {"OO:write_bool", kwlist, 0};
    PyObject *data;
    PyObject *value;
    if (unlikely(!CPyArg_ParseStackAndKeywordsSimple(args, nargs, kwnames, &parser, &data, &value))) {
        return NULL;
    }
    if (unlikely(!PyBool_Check(value))) {
        PyErr_SetString(PyExc_TypeError, "value must be a bool");
        return NULL;
    }
    if (unlikely(write_bool_internal(data, value == Py_True) == CPY_NONE_ERROR)) {
        return NULL;
    }
    Py_INCREF(Py_None);
    return Py_None;
}

/*
str format: size followed by UTF-8 bytes
    short strings (len <= 127): single byte for size as `(uint8_t)size << 1`
    long strings: \x01 followed by size as Py_ssize_t
*/

static PyObject*
read_str_internal(PyObject *data) {
    _CHECK_BUFFER(data, NULL)

    // Read string length.
    Py_ssize_t size;
    _CHECK_READ(data, 1, NULL)
    uint8_t first = _READ(data, uint8_t)
    if (likely(first != LONG_STR_TAG)) {
        // Common case: short string (len <= 127).
        size = (Py_ssize_t)(first >> 1);
    } else {
        _CHECK_READ(data, sizeof(CPyTagged), NULL)
        size = _READ(data, Py_ssize_t)
    }
    // Read string content.
    char *buf = ((BufferObject *)data)->buf;
    _CHECK_READ(data, size, NULL)
    PyObject *res = PyUnicode_FromStringAndSize(
        buf + ((BufferObject *)data)->pos, (Py_ssize_t)size
    );
    if (unlikely(res == NULL))
        return NULL;
    ((BufferObject *)data)->pos += size;
    return res;
}

static PyObject*
read_str(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames) {
    static const char * const kwlist[] = {"data", 0};
    static CPyArg_Parser parser = {"O:read_str", kwlist, 0};
    PyObject *data;
    if (unlikely(!CPyArg_ParseStackAndKeywordsOneArg(args, nargs, kwnames, &parser, &data))) {
        return NULL;
    }
    return read_str_internal(data);
}

static char
write_str_internal(PyObject *data, PyObject *value) {
    _CHECK_BUFFER(data, CPY_NONE_ERROR)

    Py_ssize_t size;
    const char *chunk = PyUnicode_AsUTF8AndSize(value, &size);
    if (unlikely(chunk == NULL))
        return CPY_NONE_ERROR;

    Py_ssize_t need;
    // Write string length.
    if (likely(size <= MAX_SHORT_LEN)) {
        // Common case: short string (len <= 127) store as single byte.
        need = size + 1;
        _CHECK_SIZE(data, need)
        _WRITE(data, uint8_t, (uint8_t)size << 1)
    } else {
        need = size + sizeof(Py_ssize_t) + 1;
        _CHECK_SIZE(data, need)
        _WRITE(data, uint8_t, LONG_STR_TAG)
        _WRITE(data, Py_ssize_t, size)
    }
    // Write string content.
    char *buf = ((BufferObject *)data)->buf;
    memcpy(buf + ((BufferObject *)data)->pos, chunk, size);
    ((BufferObject *)data)->pos += size;
    ((BufferObject *)data)->end += need;
    return CPY_NONE;
}

static PyObject*
write_str(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames) {
    static const char * const kwlist[] = {"data", "value", 0};
    static CPyArg_Parser parser = {"OO:write_str", kwlist, 0};
    PyObject *data;
    PyObject *value;
    if (unlikely(!CPyArg_ParseStackAndKeywordsSimple(args, nargs, kwnames, &parser, &data, &value))) {
        return NULL;
    }
    if (unlikely(!PyUnicode_Check(value))) {
        PyErr_SetString(PyExc_TypeError, "value must be a str");
        return NULL;
    }
    if (unlikely(write_str_internal(data, value) == CPY_NONE_ERROR)) {
        return NULL;
    }
    Py_INCREF(Py_None);
    return Py_None;
}

/*
float format:
    stored as a C double
*/

static double
read_float_internal(PyObject *data) {
    _CHECK_BUFFER(data, CPY_FLOAT_ERROR)
    _CHECK_READ(data, sizeof(double), CPY_FLOAT_ERROR)
    double res = _READ(data, double);
    return res;
}

static PyObject*
read_float(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames) {
    static const char * const kwlist[] = {"data", 0};
    static CPyArg_Parser parser = {"O:read_float", kwlist, 0};
    PyObject *data;
    if (unlikely(!CPyArg_ParseStackAndKeywordsOneArg(args, nargs, kwnames, &parser, &data))) {
        return NULL;
    }
    double retval = read_float_internal(data);
    if (unlikely(retval == CPY_FLOAT_ERROR && PyErr_Occurred())) {
        return NULL;
    }
    return PyFloat_FromDouble(retval);
}

static char
write_float_internal(PyObject *data, double value) {
    _CHECK_BUFFER(data, CPY_NONE_ERROR)
    _CHECK_SIZE(data, sizeof(double))
    _WRITE(data, double, value)
    ((BufferObject *)data)->end += sizeof(double);
    return CPY_NONE;
}

static PyObject*
write_float(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames) {
    static const char * const kwlist[] = {"data", "value", 0};
    static CPyArg_Parser parser = {"OO:write_float", kwlist, 0};
    PyObject *data;
    PyObject *value;
    if (unlikely(!CPyArg_ParseStackAndKeywordsSimple(args, nargs, kwnames, &parser, &data, &value))) {
        return NULL;
    }
    if (unlikely(!PyFloat_Check(value))) {
        PyErr_SetString(PyExc_TypeError, "value must be a float");
        return NULL;
    }
    if (unlikely(write_float_internal(data, PyFloat_AsDouble(value)) == CPY_NONE_ERROR)) {
        return NULL;
    }
    Py_INCREF(Py_None);
    return Py_None;
}

/*
int format:
    most common values (-10 <= value <= 117): single byte as `(uint8_t)(value + 10) << 1`
    medium values (fit in CPyTagged): \x01 followed by CPyTagged value
    long values (very rare): \x03 followed by decimal string (see str format)
*/

static CPyTagged
read_int_internal(PyObject *data) {
    _CHECK_BUFFER(data, CPY_INT_TAG)
    _CHECK_READ(data, 1, CPY_INT_TAG)

    uint8_t first = _READ(data, uint8_t)
    if ((first & MEDIUM_INT_TAG) == 0) {
       // Most common case: int that is small in absolute value.
       return ((Py_ssize_t)(first >> 1) + MIN_SHORT_INT) << 1;
    }
    if (first == MEDIUM_INT_TAG) {
        _CHECK_READ(data, sizeof(CPyTagged), CPY_INT_TAG)
        CPyTagged ret = _READ(data, CPyTagged)
        return ret;
    }
    // People who have literal ints not fitting in size_t should be punished :-)
    PyObject *str_ret = read_str_internal(data);
    if (unlikely(str_ret == NULL))
        return CPY_INT_TAG;
    PyObject* ret_long = PyLong_FromUnicodeObject(str_ret, 10);
    Py_DECREF(str_ret);
    return ((CPyTagged)ret_long) | CPY_INT_TAG;
}

static PyObject*
read_int(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames) {
    static const char * const kwlist[] = {"data", 0};
    static CPyArg_Parser parser = {"O:read_int", kwlist, 0};
    PyObject *data;
    if (unlikely(!CPyArg_ParseStackAndKeywordsOneArg(args, nargs, kwnames, &parser, &data))) {
        return NULL;
    }
    CPyTagged retval = read_int_internal(data);
    if (unlikely(retval == CPY_INT_TAG)) {
        return NULL;
    }
    return CPyTagged_StealAsObject(retval);
}

static char
write_int_internal(PyObject *data, CPyTagged value) {
    _CHECK_BUFFER(data, CPY_NONE_ERROR)

    if (likely((value & CPY_INT_TAG) == 0)) {
        Py_ssize_t real_value = CPyTagged_ShortAsSsize_t(value);
        if (real_value >= MIN_SHORT_INT && real_value <= MAX_SHORT_INT) {
            // Most common case: int that is small in absolute value.
            _CHECK_SIZE(data, 1)
            _WRITE(data, uint8_t, (uint8_t)(real_value - MIN_SHORT_INT) << 1)
            ((BufferObject *)data)->end += 1;
        } else {
            _CHECK_SIZE(data, sizeof(CPyTagged) + 1)
            _WRITE(data, uint8_t, MEDIUM_INT_TAG)
            _WRITE(data, CPyTagged, value)
            ((BufferObject *)data)->end += sizeof(CPyTagged) + 1;
        }
    } else {
        _CHECK_SIZE(data, 1)
        _WRITE(data, uint8_t, LONG_INT_TAG)
        ((BufferObject *)data)->end += 1;
        PyObject *str_value = PyObject_Str(CPyTagged_LongAsObject(value));
        if (unlikely(str_value == NULL))
            return CPY_NONE_ERROR;
        char res = write_str_internal(data, str_value);
        Py_DECREF(str_value);
        if (unlikely(res == CPY_NONE_ERROR))
            return CPY_NONE_ERROR;
    }
    return CPY_NONE;
}

static PyObject*
write_int(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames) {
    static const char * const kwlist[] = {"data", "value", 0};
    static CPyArg_Parser parser = {"OO:write_int", kwlist, 0};
    PyObject *data;
    PyObject *value;
    if (unlikely(!CPyArg_ParseStackAndKeywordsSimple(args, nargs, kwnames, &parser, &data, &value))) {
        return NULL;
    }
    if (unlikely(!PyLong_Check(value))) {
        PyErr_SetString(PyExc_TypeError, "value must be an int");
        return NULL;
    }
    CPyTagged tagged_value = CPyTagged_BorrowFromObject(value);
    if (unlikely(write_int_internal(data, tagged_value) == CPY_NONE_ERROR)) {
        return NULL;
    }
    Py_INCREF(Py_None);
    return Py_None;
}

/*
integer tag format (0 <= t <= 255):
    stored as a uint8_t
*/

static uint8_t
read_tag_internal(PyObject *data) {
    _CHECK_BUFFER(data, CPY_LL_UINT_ERROR)
    _CHECK_READ(data, 1, CPY_LL_UINT_ERROR)
    uint8_t ret = _READ(data, uint8_t)
    return ret;
}

static PyObject*
read_tag(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames) {
    static const char * const kwlist[] = {"data", 0};
    static CPyArg_Parser parser = {"O:read_tag", kwlist, 0};
    PyObject *data;
    if (unlikely(!CPyArg_ParseStackAndKeywordsOneArg(args, nargs, kwnames, &parser, &data))) {
        return NULL;
    }
    uint8_t retval = read_tag_internal(data);
    if (unlikely(retval == CPY_LL_UINT_ERROR && PyErr_Occurred())) {
        return NULL;
    }
    return PyLong_FromLong(retval);
}

static char
write_tag_internal(PyObject *data, uint8_t value) {
    _CHECK_BUFFER(data, CPY_NONE_ERROR)
    _CHECK_SIZE(data, 1)
    _WRITE(data, uint8_t, value)
    ((BufferObject *)data)->end += 1;
    return CPY_NONE;
}

static PyObject*
write_tag(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames) {
    static const char * const kwlist[] = {"data", "value", 0};
    static CPyArg_Parser parser = {"OO:write_tag", kwlist, 0};
    PyObject *data;
    PyObject *value;
    if (unlikely(!CPyArg_ParseStackAndKeywordsSimple(args, nargs, kwnames, &parser, &data, &value))) {
        return NULL;
    }
    uint8_t unboxed = CPyLong_AsUInt8(value);
    if (unlikely(unboxed == CPY_LL_UINT_ERROR && PyErr_Occurred())) {
        CPy_TypeError("u8", value);
        return NULL;
    }
    if (unlikely(write_tag_internal(data, unboxed) == CPY_NONE_ERROR)) {
        return NULL;
    }
    Py_INCREF(Py_None);
    return Py_None;
}

static PyMethodDef native_internal_module_methods[] = {
    {"write_bool", (PyCFunction)write_bool, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("write a bool")},
    {"read_bool", (PyCFunction)read_bool, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("read a bool")},
    {"write_str", (PyCFunction)write_str, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("write a string")},
    {"read_str", (PyCFunction)read_str, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("read a string")},
    {"write_float", (PyCFunction)write_float, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("write a float")},
    {"read_float", (PyCFunction)read_float, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("read a float")},
    {"write_int", (PyCFunction)write_int, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("write an int")},
    {"read_int", (PyCFunction)read_int, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("read an int")},
    {"write_tag", (PyCFunction)write_tag, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("write a short int")},
    {"read_tag", (PyCFunction)read_tag, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("read a short int")},
    {NULL, NULL, 0, NULL}
};

static int
NativeInternal_ABI_Version(void) {
    return NATIVE_INTERNAL_ABI_VERSION;
}

static int
native_internal_module_exec(PyObject *m)
{
    if (PyType_Ready(&BufferType) < 0) {
        return -1;
    }
    if (PyModule_AddObjectRef(m, "Buffer", (PyObject *) &BufferType) < 0) {
        return -1;
    }

    // Export mypy internal C API, be careful with the order!
    static void *NativeInternal_API[14] = {
        (void *)Buffer_internal,
        (void *)Buffer_internal_empty,
        (void *)Buffer_getvalue_internal,
        (void *)write_bool_internal,
        (void *)read_bool_internal,
        (void *)write_str_internal,
        (void *)read_str_internal,
        (void *)write_float_internal,
        (void *)read_float_internal,
        (void *)write_int_internal,
        (void *)read_int_internal,
        (void *)write_tag_internal,
        (void *)read_tag_internal,
        (void *)NativeInternal_ABI_Version,
    };
    PyObject *c_api_object = PyCapsule_New((void *)NativeInternal_API, "native_internal._C_API", NULL);
    if (PyModule_Add(m, "_C_API", c_api_object) < 0) {
        return -1;
    }
    return 0;
}

static PyModuleDef_Slot native_internal_module_slots[] = {
    {Py_mod_exec, native_internal_module_exec},
#ifdef Py_MOD_GIL_NOT_USED
    {Py_mod_gil, Py_MOD_GIL_NOT_USED},
#endif
    {0, NULL}
};

static PyModuleDef native_internal_module = {
    .m_base = PyModuleDef_HEAD_INIT,
    .m_name = "native_internal",
    .m_doc = "Mypy cache serialization utils",
    .m_size = 0,
    .m_methods = native_internal_module_methods,
    .m_slots = native_internal_module_slots,
};

PyMODINIT_FUNC
PyInit_native_internal(void)
{
    return PyModuleDef_Init(&native_internal_module);
}