/**
 * Copyright (c) 2018-present, Gregory Szorc
 * All rights reserved.
 *
 * This software may be modified and distributed under the terms
 * of the BSD license. See the LICENSE file for details.
 */

#include "python-zstandard.h"

extern PyObject *ZstdError;

static void
ZstdCompressionChunkerIterator_dealloc(ZstdCompressionChunkerIterator *self) {
    Py_XDECREF(self->chunker);

    PyObject_Del(self);
}

static PyObject *ZstdCompressionChunkerIterator_iter(PyObject *self) {
    Py_INCREF(self);
    return self;
}

static PyObject *
ZstdCompressionChunkerIterator_iternext(ZstdCompressionChunkerIterator *self) {
    size_t zresult;
    PyObject *chunk;
    ZstdCompressionChunker *chunker = self->chunker;
    ZSTD_EndDirective zFlushMode;

    if (self->mode != compressionchunker_mode_normal &&
        chunker->input.pos != chunker->input.size) {
        PyErr_SetString(ZstdError, "input should have been fully consumed "
                                   "before calling flush() or finish()");
        return NULL;
    }

    if (chunker->finished) {
        return NULL;
    }

    /* If we have data left in the input, consume it. */
    while (chunker->input.pos < chunker->input.size) {
        Py_BEGIN_ALLOW_THREADS zresult =
            ZSTD_compressStream2(chunker->compressor->cctx, &chunker->output,
                                 &chunker->input, ZSTD_e_continue);
        Py_END_ALLOW_THREADS

            /* Input is fully consumed. */
            if (chunker->input.pos == chunker->input.size) {
            chunker->input.src = NULL;
            chunker->input.pos = 0;
            chunker->input.size = 0;
            PyBuffer_Release(&chunker->inBuffer);
        }

        if (ZSTD_isError(zresult)) {
            PyErr_Format(ZstdError, "zstd compress error: %s",
                         ZSTD_getErrorName(zresult));
            return NULL;
        }

        /* If it produced a full output chunk, emit it. */
        if (chunker->output.pos == chunker->output.size) {
            chunk = PyBytes_FromStringAndSize(chunker->output.dst,
                                              chunker->output.pos);
            if (!chunk) {
                return NULL;
            }

            chunker->output.pos = 0;

            return chunk;
        }

        /* Else continue to compress available input data. */
    }

    /* We also need this here for the special case of an empty input buffer. */
    if (chunker->input.pos == chunker->input.size) {
        chunker->input.src = NULL;
        chunker->input.pos = 0;
        chunker->input.size = 0;
        PyBuffer_Release(&chunker->inBuffer);
    }

    /* No more input data. A partial chunk may be in chunker->output.
     * If we're in normal compression mode, we're done. Otherwise if we're in
     * flush or finish mode, we need to emit what data remains.
     */
    if (self->mode == compressionchunker_mode_normal) {
        /* We don't need to set StopIteration. */
        return NULL;
    }

    if (self->mode == compressionchunker_mode_flush) {
        zFlushMode = ZSTD_e_flush;
    }
    else if (self->mode == compressionchunker_mode_finish) {
        zFlushMode = ZSTD_e_end;
    }
    else {
        PyErr_SetString(ZstdError,
                        "unhandled compression mode; this should never happen");
        return NULL;
    }

    Py_BEGIN_ALLOW_THREADS zresult =
        ZSTD_compressStream2(chunker->compressor->cctx, &chunker->output,
                             &chunker->input, zFlushMode);
    Py_END_ALLOW_THREADS

        if (ZSTD_isError(zresult)) {
        PyErr_Format(ZstdError, "zstd compress error: %s",
                     ZSTD_getErrorName(zresult));
        return NULL;
    }

    if (!zresult && chunker->output.pos == 0) {
        return NULL;
    }

    chunk = PyBytes_FromStringAndSize(chunker->output.dst, chunker->output.pos);
    if (!chunk) {
        return NULL;
    }

    chunker->output.pos = 0;

    if (!zresult && self->mode == compressionchunker_mode_finish) {
        chunker->finished = 1;
    }

    return chunk;
}

PyTypeObject ZstdCompressionChunkerIteratorType = {
    PyVarObject_HEAD_INIT(
        NULL, 0) "zstd.ZstdCompressionChunkerIterator", /* tp_name */
    sizeof(ZstdCompressionChunkerIterator),             /* tp_basicsize */
    0,                                                  /* tp_itemsize */
    (destructor)ZstdCompressionChunkerIterator_dealloc, /* tp_dealloc */
    0,                                                  /* tp_print */
    0,                                                  /* tp_getattr */
    0,                                                  /* tp_setattr */
    0,                                                  /* tp_compare */
    0,                                                  /* tp_repr */
    0,                                                  /* tp_as_number */
    0,                                                  /* tp_as_sequence */
    0,                                                  /* tp_as_mapping */
    0,                                                  /* tp_hash */
    0,                                                  /* tp_call */
    0,                                                  /* tp_str */
    0,                                                  /* tp_getattro */
    0,                                                  /* tp_setattro */
    0,                                                  /* tp_as_buffer */
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,           /* tp_flags */
    0,                                                  /* tp_doc */
    0,                                                  /* tp_traverse */
    0,                                                  /* tp_clear */
    0,                                                  /* tp_richcompare */
    0,                                                  /* tp_weaklistoffset */
    ZstdCompressionChunkerIterator_iter,                /* tp_iter */
    (iternextfunc)ZstdCompressionChunkerIterator_iternext, /* tp_iternext */
    0,                                                     /* tp_methods */
    0,                                                     /* tp_members */
    0,                                                     /* tp_getset */
    0,                                                     /* tp_base */
    0,                                                     /* tp_dict */
    0,                                                     /* tp_descr_get */
    0,                                                     /* tp_descr_set */
    0,                                                     /* tp_dictoffset */
    0,                                                     /* tp_init */
    0,                                                     /* tp_alloc */
    PyType_GenericNew,                                     /* tp_new */
};

static void ZstdCompressionChunker_dealloc(ZstdCompressionChunker *self) {
    PyBuffer_Release(&self->inBuffer);
    self->input.src = NULL;

    PyMem_Free(self->output.dst);
    self->output.dst = NULL;

    Py_XDECREF(self->compressor);

    PyObject_Del(self);
}

static ZstdCompressionChunkerIterator *
ZstdCompressionChunker_compress(ZstdCompressionChunker *self, PyObject *args,
                                PyObject *kwargs) {
    static char *kwlist[] = {"data", NULL};

    ZstdCompressionChunkerIterator *result;

    if (self->finished) {
        PyErr_SetString(ZstdError,
                        "cannot call compress() after compression finished");
        return NULL;
    }

    if (self->inBuffer.obj) {
        PyErr_SetString(ZstdError, "cannot perform operation before consuming "
                                   "output from previous operation");
        return NULL;
    }

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*:compress", kwlist,
                                     &self->inBuffer)) {
        return NULL;
    }

    result = (ZstdCompressionChunkerIterator *)PyObject_CallObject(
        (PyObject *)&ZstdCompressionChunkerIteratorType, NULL);
    if (!result) {
        PyBuffer_Release(&self->inBuffer);
        return NULL;
    }

    self->input.src = self->inBuffer.buf;
    self->input.size = self->inBuffer.len;
    self->input.pos = 0;

    result->chunker = self;
    Py_INCREF(result->chunker);

    result->mode = compressionchunker_mode_normal;

    return result;
}

static ZstdCompressionChunkerIterator *
ZstdCompressionChunker_finish(ZstdCompressionChunker *self) {
    ZstdCompressionChunkerIterator *result;

    if (self->finished) {
        PyErr_SetString(ZstdError,
                        "cannot call finish() after compression finished");
        return NULL;
    }

    if (self->inBuffer.obj) {
        PyErr_SetString(ZstdError, "cannot call finish() before consuming "
                                   "output from previous operation");
        return NULL;
    }

    result = (ZstdCompressionChunkerIterator *)PyObject_CallObject(
        (PyObject *)&ZstdCompressionChunkerIteratorType, NULL);
    if (!result) {
        return NULL;
    }

    result->chunker = self;
    Py_INCREF(result->chunker);

    result->mode = compressionchunker_mode_finish;

    return result;
}

static ZstdCompressionChunkerIterator *
ZstdCompressionChunker_flush(ZstdCompressionChunker *self, PyObject *args,
                             PyObject *kwargs) {
    ZstdCompressionChunkerIterator *result;

    if (self->finished) {
        PyErr_SetString(ZstdError,
                        "cannot call flush() after compression finished");
        return NULL;
    }

    if (self->inBuffer.obj) {
        PyErr_SetString(ZstdError, "cannot call flush() before consuming "
                                   "output from previous operation");
        return NULL;
    }

    result = (ZstdCompressionChunkerIterator *)PyObject_CallObject(
        (PyObject *)&ZstdCompressionChunkerIteratorType, NULL);
    if (!result) {
        return NULL;
    }

    result->chunker = self;
    Py_INCREF(result->chunker);

    result->mode = compressionchunker_mode_flush;

    return result;
}

static PyMethodDef ZstdCompressionChunker_methods[] = {
    {"compress", (PyCFunction)ZstdCompressionChunker_compress,
     METH_VARARGS | METH_KEYWORDS, PyDoc_STR("compress data")},
    {"finish", (PyCFunction)ZstdCompressionChunker_finish, METH_NOARGS,
     PyDoc_STR("finish compression operation")},
    {"flush", (PyCFunction)ZstdCompressionChunker_flush,
     METH_VARARGS | METH_KEYWORDS, PyDoc_STR("finish compression operation")},
    {NULL, NULL}};

PyTypeObject ZstdCompressionChunkerType = {
    PyVarObject_HEAD_INIT(NULL,
                          0) "zstd.ZstdCompressionChunkerType", /* tp_name */
    sizeof(ZstdCompressionChunker),             /* tp_basicsize */
    0,                                          /* tp_itemsize */
    (destructor)ZstdCompressionChunker_dealloc, /* tp_dealloc */
    0,                                          /* tp_print */
    0,                                          /* tp_getattr */
    0,                                          /* tp_setattr */
    0,                                          /* tp_compare */
    0,                                          /* tp_repr */
    0,                                          /* tp_as_number */
    0,                                          /* tp_as_sequence */
    0,                                          /* tp_as_mapping */
    0,                                          /* tp_hash */
    0,                                          /* tp_call */
    0,                                          /* tp_str */
    0,                                          /* tp_getattro */
    0,                                          /* tp_setattro */
    0,                                          /* tp_as_buffer */
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,   /* tp_flags */
    0,                                          /* tp_doc */
    0,                                          /* tp_traverse */
    0,                                          /* tp_clear */
    0,                                          /* tp_richcompare */
    0,                                          /* tp_weaklistoffset */
    0,                                          /* tp_iter */
    0,                                          /* tp_iternext */
    ZstdCompressionChunker_methods,             /* tp_methods */
    0,                                          /* tp_members */
    0,                                          /* tp_getset */
    0,                                          /* tp_base */
    0,                                          /* tp_dict */
    0,                                          /* tp_descr_get */
    0,                                          /* tp_descr_set */
    0,                                          /* tp_dictoffset */
    0,                                          /* tp_init */
    0,                                          /* tp_alloc */
    PyType_GenericNew,                          /* tp_new */
};

void compressionchunker_module_init(PyObject *module) {
    Py_SET_TYPE(&ZstdCompressionChunkerIteratorType, &PyType_Type);
    if (PyType_Ready(&ZstdCompressionChunkerIteratorType) < 0) {
        return;
    }

    Py_SET_TYPE(&ZstdCompressionChunkerType, &PyType_Type);
    if (PyType_Ready(&ZstdCompressionChunkerType) < 0) {
        return;
    }
}
