/**
 * Copyright (c) 2016-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"

#define min(a, b) (((a) < (b)) ? (a) : (b))

extern PyObject *ZstdError;

static void ZstdCompressorIterator_dealloc(ZstdCompressorIterator *self) {
    Py_XDECREF(self->readResult);
    Py_XDECREF(self->compressor);
    Py_XDECREF(self->reader);

    if (self->buffer.buf) {
        PyBuffer_Release(&self->buffer);
        memset(&self->buffer, 0, sizeof(self->buffer));
    }

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

    PyObject_Del(self);
}

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

static PyObject *ZstdCompressorIterator_iternext(ZstdCompressorIterator *self) {
    size_t zresult;
    PyObject *readResult = NULL;
    PyObject *chunk;
    char *readBuffer;
    Py_ssize_t readSize = 0;
    Py_ssize_t bufferRemaining;

    if (self->finishedOutput) {
        PyErr_SetString(PyExc_StopIteration, "output flushed");
        return NULL;
    }

feedcompressor:

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

            /* Release the Python object holding the input buffer. */
            if (self->input.pos == self->input.size) {
            self->input.src = NULL;
            self->input.pos = 0;
            self->input.size = 0;
            Py_DECREF(self->readResult);
            self->readResult = NULL;
        }

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

        /* If it produced output data, emit it. */
        if (self->output.pos) {
            chunk =
                PyBytes_FromStringAndSize(self->output.dst, self->output.pos);
            self->output.pos = 0;
            return chunk;
        }
    }

    /* We should never have output data sitting around after a previous call. */
    assert(self->output.pos == 0);

    /* The code above should have either emitted a chunk and returned or
    consumed the entire input buffer. So the state of the input buffer is not
    relevant. */
    if (!self->finishedInput) {
        if (self->reader) {
            readResult =
                PyObject_CallMethod(self->reader, "read", "I", self->inSize);
            if (!readResult) {
                return NULL;
            }

            PyBytes_AsStringAndSize(readResult, &readBuffer, &readSize);
        }
        else {
            assert(self->buffer.buf);

            /* Only support contiguous C arrays. */
            assert(self->buffer.strides == NULL &&
                   self->buffer.suboffsets == NULL);
            assert(self->buffer.itemsize == 1);

            readBuffer = (char *)self->buffer.buf + self->bufferOffset;
            bufferRemaining = self->buffer.len - self->bufferOffset;
            readSize = min(bufferRemaining, (Py_ssize_t)self->inSize);
            self->bufferOffset += readSize;
        }

        if (0 == readSize) {
            Py_XDECREF(readResult);
            self->finishedInput = 1;
        }
        else {
            self->readResult = readResult;
        }
    }

    /* EOF */
    if (0 == readSize) {
        self->input.src = NULL;
        self->input.size = 0;
        self->input.pos = 0;

        zresult = ZSTD_compressStream2(self->compressor->cctx, &self->output,
                                       &self->input, ZSTD_e_end);
        if (ZSTD_isError(zresult)) {
            PyErr_Format(ZstdError, "error ending compression stream: %s",
                         ZSTD_getErrorName(zresult));
            return NULL;
        }

        assert(self->output.pos);

        if (0 == zresult) {
            self->finishedOutput = 1;
        }

        chunk = PyBytes_FromStringAndSize(self->output.dst, self->output.pos);
        self->output.pos = 0;
        return chunk;
    }

    /* New data from reader. Feed into compressor. */
    self->input.src = readBuffer;
    self->input.size = readSize;
    self->input.pos = 0;

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

        /* The input buffer currently points to memory managed by Python
        (readBuffer). This object was allocated by this function. If it wasn't
        fully consumed, we need to release it in a subsequent function call.
        If it is fully consumed, do that now.
        */
        if (self->input.pos == self->input.size) {
        self->input.src = NULL;
        self->input.pos = 0;
        self->input.size = 0;
        Py_XDECREF(self->readResult);
        self->readResult = NULL;
    }

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

    assert(self->input.pos <= self->input.size);

    /* If we didn't write anything, start the process over. */
    if (0 == self->output.pos) {
        goto feedcompressor;
    }

    chunk = PyBytes_FromStringAndSize(self->output.dst, self->output.pos);
    self->output.pos = 0;
    return chunk;
}

PyTypeObject ZstdCompressorIteratorType = {
    PyVarObject_HEAD_INIT(NULL, 0) "zstd.ZstdCompressorIterator", /* tp_name */
    sizeof(ZstdCompressorIterator),                /* tp_basicsize */
    0,                                             /* tp_itemsize */
    (destructor)ZstdCompressorIterator_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 */
    ZstdCompressorIterator_iter,                   /* tp_iter */
    (iternextfunc)ZstdCompressorIterator_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 */
};

void compressoriterator_module_init(PyObject *mod) {
    Py_SET_TYPE(&ZstdCompressorIteratorType, &PyType_Type);
    if (PyType_Ready(&ZstdCompressorIteratorType) < 0) {
        return;
    }
}
