/* -*- mode: c++; c-basic-offset: 4 -*- */

/* For linux, png.h must be imported before Python.h because
   png.h needs to be the one to define setjmp.
   Undefining _POSIX_C_SOURCE and _XOPEN_SOURCE stops a couple
   of harmless warnings.
*/
#define PY_SSIZE_T_CLEAN

extern "C" {
#   include <png.h>
#   ifdef _POSIX_C_SOURCE
#       undef _POSIX_C_SOURCE
#   endif
#   ifndef _AIX
#   ifdef _XOPEN_SOURCE
#       undef _XOPEN_SOURCE
#   endif
#   endif
}

#include "numpy_cpp.h"

#   include <vector>
#   include "Python.h"


// As reported in [3082058] build _png.so on aix
#ifdef _AIX
#undef jmpbuf
#endif

struct buffer_t {
    PyObject *str;
    size_t cursor;
    size_t size;
};


static void write_png_data_buffer(png_structp png_ptr, png_bytep data, png_size_t length)
{
    buffer_t *buff = (buffer_t *)png_get_io_ptr(png_ptr);
    if (buff->cursor + length < buff->size) {
        memcpy(PyBytes_AS_STRING(buff->str) + buff->cursor, data, length);
        buff->cursor += length;
    }
}

static void flush_png_data_buffer(png_structp png_ptr)
{

}

static void write_png_data(png_structp png_ptr, png_bytep data, png_size_t length)
{
    PyObject *py_file_obj = (PyObject *)png_get_io_ptr(png_ptr);
    PyObject *write_method = PyObject_GetAttrString(py_file_obj, "write");
    PyObject *result = NULL;
    if (write_method) {
        result = PyObject_CallFunction(write_method, (char *)"y#", data, length);
    }
    Py_XDECREF(write_method);
    Py_XDECREF(result);
}

static void flush_png_data(png_structp png_ptr)
{
    PyObject *py_file_obj = (PyObject *)png_get_io_ptr(png_ptr);
    PyObject *flush_method = PyObject_GetAttrString(py_file_obj, "flush");
    PyObject *result = NULL;
    if (flush_method) {
        result = PyObject_CallFunction(flush_method, (char *)"");
    }
    Py_XDECREF(flush_method);
    Py_XDECREF(result);
}

const char *Py_write_png__doc__ =
    "write_png(buffer, file, dpi=0, compression=6, filter=auto, metadata=None)\n"
    "\n"
    "Parameters\n"
    "----------\n"
    "buffer : numpy array of image data\n"
    "    Must be an MxNxD array of dtype uint8.\n"
    "    - if D is 1, the image is greyscale\n"
    "    - if D is 3, the image is RGB\n"
    "    - if D is 4, the image is RGBA\n"
    "\n"
    "file : binary-mode file-like object or None\n"
    "    If None, a byte string containing the PNG data will be returned\n"
    "\n"
    "dpi : float\n"
    "    The dpi to store in the file metadata.\n"
    "\n"
    "compression : int\n"
    "    The level of lossless zlib compression to apply.  0 indicates no\n"
    "    compression.  Values 1-9 indicate low/fast through high/slow\n"
    "    compression.  Default is 6.\n"
    "\n"
    "filter : int\n"
    "    Filter to apply.  Must be one of the constants: PNG_FILTER_NONE,\n"
    "    PNG_FILTER_SUB, PNG_FILTER_UP, PNG_FILTER_AVG, PNG_FILTER_PAETH.\n"
    "    See the PNG standard for more information.\n"
    "    If not provided, libpng will try to automatically determine the\n"
    "    best filter on a line-by-line basis.\n"
    "\n"
    "metadata : dictionary\n"
    "    The keyword-text pairs that are stored as comments in the image.\n"
    "    Keys must be shorter than 79 chars. The only supported encoding\n"
    "    for both keywords and values is Latin-1 (ISO 8859-1).\n"
    "    Examples given in the PNG Specification are:\n"
    "    - Title: Short (one line) title or caption for image\n"
    "    - Author: Name of image's creator\n"
    "    - Description: Description of image (possibly long)\n"
    "    - Copyright: Copyright notice\n"
    "    - Creation Time: Time of original image creation\n"
    "                     (usually RFC 1123 format, see below)\n"
    "    - Software: Software used to create the image\n"
    "    - Disclaimer: Legal disclaimer\n"
    "    - Warning: Warning of nature of content\n"
    "    - Source: Device used to create the image\n"
    "    - Comment: Miscellaneous comment; conversion\n"
    "               from other image format\n"
    "\n"
    "    For more details see the PNG specification:\n"
    "    https://www.w3.org/TR/2003/REC-PNG-20031110/#11keywords\n"
    "\n"
    "Returns\n"
    "-------\n"
    "buffer : bytes or None\n"
    "    Byte string containing the PNG content if None was passed in for\n"
    "    file, otherwise None is returned.\n";

// this code is heavily adapted from
// https://www.object-craft.com.au/projects/paint/ which licensed under the
// (BSD compatible) LICENSE_PAINT which is included in this distribution.
static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
{
    numpy::array_view<unsigned char, 3> buffer;
    PyObject *filein;
    PyObject *metadata = NULL;
    PyObject *meta_key, *meta_val;
    png_text *text;
    Py_ssize_t pos = 0;
    int meta_pos = 0;
    Py_ssize_t meta_size;
    double dpi = 0;
    int compression = 6;
    int filter = -1;
    const char *names[] = { "buffer", "file", "dpi", "compression", "filter", "metadata", NULL };

    // We don't need strict contiguity, just for each row to be
    // contiguous, and libpng has special handling for getting RGB out
    // of RGBA, ARGB or BGR. But the simplest thing to do is to
    // enforce contiguity using array_view::converter_contiguous.
    if (!PyArg_ParseTupleAndKeywords(args,
                                     kwds,
                                     "O&O|diiO:write_png",
                                     (char **)names,
                                     &buffer.converter_contiguous,
                                     &buffer,
                                     &filein,
                                     &dpi,
                                     &compression,
                                     &filter,
                                     &metadata)) {
        return NULL;
    }

    png_uint_32 width = (png_uint_32)buffer.dim(1);
    png_uint_32 height = (png_uint_32)buffer.dim(0);
    npy_intp channels = buffer.dim(2);
    std::vector<png_bytep> row_pointers(height);
    for (png_uint_32 row = 0; row < (png_uint_32)height; ++row) {
        row_pointers[row] = (png_bytep)&buffer(row, 0, 0);
    }

    png_structp png_ptr = NULL;
    png_infop info_ptr = NULL;
    struct png_color_8_struct sig_bit;
    int png_color_type;
    buffer_t buff;
    buff.str = NULL;

    switch (channels) {
    case 1:
        png_color_type = PNG_COLOR_TYPE_GRAY;
        break;
    case 3:
        png_color_type = PNG_COLOR_TYPE_RGB;
        break;
    case 4:
        png_color_type = PNG_COLOR_TYPE_RGB_ALPHA;
        break;
    default:
        PyErr_SetString(PyExc_ValueError,
                        "Buffer must be an NxMxD array with D in 1, 3, 4 "
                        "(grayscale, RGB, RGBA)");
        goto exit;
    }

    if (compression < 0 || compression > 9) {
        PyErr_Format(PyExc_ValueError,
                     "compression must be in range 0-9, got %d", compression);
        goto exit;
    }

    if (filein == Py_None) {
        buff.size = width * height * 4 + 1024;
        buff.str = PyBytes_FromStringAndSize(NULL, buff.size);
        if (buff.str == NULL) {
            goto exit;
        }
        buff.cursor = 0;
    } else {
        PyErr_Clear();
        PyObject *write_method = PyObject_GetAttrString(filein, "write");
        if (!(write_method && PyCallable_Check(write_method))) {
            Py_XDECREF(write_method);
            PyErr_SetString(PyExc_TypeError,
                            "Object does not appear to be a file-like object");
            goto exit;
        }
        Py_XDECREF(write_method);
    }

    png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
    if (png_ptr == NULL) {
        PyErr_SetString(PyExc_RuntimeError, "Could not create write struct");
        goto exit;
    }

    png_set_compression_level(png_ptr, compression);
    if (filter >= 0) {
        png_set_filter(png_ptr, 0, filter);
    }

    info_ptr = png_create_info_struct(png_ptr);
    if (info_ptr == NULL) {
        PyErr_SetString(PyExc_RuntimeError, "Could not create info struct");
        goto exit;
    }

    if (setjmp(png_jmpbuf(png_ptr))) {
        PyErr_SetString(PyExc_RuntimeError, "libpng signaled error");
        goto exit;
    }

    if (buff.str) {
        png_set_write_fn(png_ptr, (void *)&buff, &write_png_data_buffer, &flush_png_data_buffer);
    } else {
        png_set_write_fn(png_ptr, (void *)filein, &write_png_data, &flush_png_data);
    }
    png_set_IHDR(png_ptr,
                 info_ptr,
                 width,
                 height,
                 8,
                 png_color_type,
                 PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_BASE,
                 PNG_FILTER_TYPE_BASE);

    // Save the dpi of the image in the file
    if (dpi > 0.0) {
        png_uint_32 dots_per_meter = (png_uint_32)(dpi / (2.54 / 100.0));
        png_set_pHYs(png_ptr, info_ptr, dots_per_meter, dots_per_meter, PNG_RESOLUTION_METER);
    }

#ifdef PNG_TEXT_SUPPORTED
    // Save the metadata
    if (metadata != NULL && metadata != Py_None) {
        if (!PyDict_Check(metadata)) {
            PyErr_SetString(PyExc_TypeError, "metadata must be a dict or None");
            goto exit;
        }
        meta_size = PyDict_Size(metadata);
        text = new png_text[meta_size];

        while (PyDict_Next(metadata, &pos, &meta_key, &meta_val)) {
            text[meta_pos].compression = PNG_TEXT_COMPRESSION_NONE;
            if (PyUnicode_Check(meta_key)) {
                PyObject *temp_key = PyUnicode_AsEncodedString(meta_key, "latin_1", "strict");
                if (temp_key != NULL) {
                    text[meta_pos].key = PyBytes_AsString(temp_key);
                }
            } else if (PyBytes_Check(meta_key)) {
                text[meta_pos].key = PyBytes_AsString(meta_key);
            } else {
                char invalid_key[79];
                sprintf(invalid_key,"INVALID KEY %d", meta_pos);
                text[meta_pos].key = invalid_key;
            }
            if (PyUnicode_Check(meta_val)) {
                PyObject *temp_val = PyUnicode_AsEncodedString(meta_val, "latin_1", "strict");
                if (temp_val != NULL) {
                    text[meta_pos].text = PyBytes_AsString(temp_val);
                }
            } else if (PyBytes_Check(meta_val)) {
                text[meta_pos].text = PyBytes_AsString(meta_val);
            } else {
                text[meta_pos].text = (char *)"Invalid value in metadata";
            }
#ifdef PNG_iTXt_SUPPORTED
            text[meta_pos].lang = NULL;
#endif
            meta_pos++;
        }
        png_set_text(png_ptr, info_ptr, text, meta_size);
        delete[] text;
    }
#endif

    sig_bit.alpha = 0;
    switch (png_color_type) {
    case PNG_COLOR_TYPE_GRAY:
        sig_bit.gray = 8;
        sig_bit.red = 0;
        sig_bit.green = 0;
        sig_bit.blue = 0;
        break;
    case PNG_COLOR_TYPE_RGB_ALPHA:
        sig_bit.alpha = 8;
        // fall through
    case PNG_COLOR_TYPE_RGB:
        sig_bit.gray = 0;
        sig_bit.red = 8;
        sig_bit.green = 8;
        sig_bit.blue = 8;
        break;
    default:
        PyErr_SetString(PyExc_RuntimeError, "internal error, bad png_color_type");
        goto exit;
    }
    png_set_sBIT(png_ptr, info_ptr, &sig_bit);

    png_write_info(png_ptr, info_ptr);
    png_write_image(png_ptr, &row_pointers[0]);
    png_write_end(png_ptr, info_ptr);

exit:
    if (png_ptr && info_ptr) {
        png_destroy_write_struct(&png_ptr, &info_ptr);
    }
    if (PyErr_Occurred()) {
        Py_XDECREF(buff.str);
        return NULL;
    } else {
        if (buff.str) {
            _PyBytes_Resize(&buff.str, buff.cursor);
            return buff.str;
        }
        Py_RETURN_NONE;
    }
}

static void _read_png_data(PyObject *py_file_obj, png_bytep data, png_size_t length)
{
    PyObject *read_method = PyObject_GetAttrString(py_file_obj, "read");
    PyObject *result = NULL;
    char *buffer;
    Py_ssize_t bufflen;
    if (read_method) {
        result = PyObject_CallFunction(read_method, (char *)"i", length);
        if (result) {
            if (PyBytes_AsStringAndSize(result, &buffer, &bufflen) == 0) {
                if (bufflen == (Py_ssize_t)length) {
                    memcpy(data, buffer, length);
                } else {
                    PyErr_SetString(PyExc_IOError, "read past end of file");
                }
            } else {
                PyErr_SetString(PyExc_IOError, "failed to copy buffer");
            }
        } else  {
            PyErr_SetString(PyExc_IOError, "failed to read file");
        }


    }
    Py_XDECREF(read_method);
    Py_XDECREF(result);
}

static void read_png_data(png_structp png_ptr, png_bytep data, png_size_t length)
{
    PyObject *py_file_obj = (PyObject *)png_get_io_ptr(png_ptr);
    _read_png_data(py_file_obj, data, length);
    if (PyErr_Occurred()) {
        png_error(png_ptr, "failed to read file");
    }

}

static PyObject *_read_png(PyObject *filein, bool float_result)
{
    png_byte header[8]; // 8 is the maximum size that can be checked
    png_structp png_ptr = NULL;
    png_infop info_ptr = NULL;
    int num_dims;
    std::vector<png_bytep> row_pointers;
    png_uint_32 width = 0;
    png_uint_32 height = 0;
    int bit_depth;
    PyObject *result = NULL;

    // TODO: Remove direct calls to Numpy API here
    PyErr_Clear();

    PyObject *read_method = PyObject_GetAttrString(filein, "read");
    if (!(read_method && PyCallable_Check(read_method))) {
        Py_XDECREF(read_method);
        PyErr_SetString(PyExc_TypeError,
                        "Object does not appear to be a file-like object");
        goto exit;
    }
    Py_XDECREF(read_method);
    _read_png_data(filein, header, 8);
    if (PyErr_Occurred()) {
        goto exit;
    }

    if (png_sig_cmp(header, 0, 8)) {
        PyErr_SetString(PyExc_ValueError, "invalid PNG header");
        goto exit;
    }

    /* initialize stuff */
    png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);

    if (!png_ptr) {
        PyErr_SetString(PyExc_RuntimeError, "png_create_read_struct failed");
        goto exit;
    }

    info_ptr = png_create_info_struct(png_ptr);
    if (!info_ptr) {
        PyErr_SetString(PyExc_RuntimeError, "png_create_info_struct failed");
        goto exit;
    }

    if (setjmp(png_jmpbuf(png_ptr))) {
        if (!PyErr_Occurred()) {
            PyErr_SetString(PyExc_RuntimeError, "error setting jump");
        }
        goto exit;
    }

    png_set_read_fn(png_ptr, (void *)filein, &read_png_data);
    png_set_sig_bytes(png_ptr, 8);
    png_read_info(png_ptr, info_ptr);
    width = png_get_image_width(png_ptr, info_ptr);
    height = png_get_image_height(png_ptr, info_ptr);
    bit_depth = png_get_bit_depth(png_ptr, info_ptr);

    // Unpack 1, 2, and 4-bit images
    if (bit_depth < 8) {
        png_set_packing(png_ptr);
    }

    // If sig bits are set, shift data
    png_color_8p sig_bit;
    if ((png_get_color_type(png_ptr, info_ptr) != PNG_COLOR_TYPE_PALETTE) &&
        png_get_sBIT(png_ptr, info_ptr, &sig_bit)) {
        png_set_shift(png_ptr, sig_bit);
    }

#if NPY_BYTE_ORDER == NPY_LITTLE_ENDIAN
    // Convert big endian to little
    if (bit_depth == 16) {
        png_set_swap(png_ptr);
    }
#endif

    // Convert palletes to full RGB
    if (png_get_color_type(png_ptr, info_ptr) == PNG_COLOR_TYPE_PALETTE) {
        png_set_palette_to_rgb(png_ptr);
        bit_depth = 8;
    }

    // If there's an alpha channel convert gray to RGB
    if (png_get_color_type(png_ptr, info_ptr) == PNG_COLOR_TYPE_GRAY_ALPHA) {
        png_set_gray_to_rgb(png_ptr);
    }

    png_set_interlace_handling(png_ptr);
    png_read_update_info(png_ptr, info_ptr);

    row_pointers.resize(height);
    for (png_uint_32 row = 0; row < height; row++) {
        row_pointers[row] = new png_byte[png_get_rowbytes(png_ptr, info_ptr)];
    }

    png_read_image(png_ptr, &row_pointers[0]);

    npy_intp dimensions[3];
    dimensions[0] = height; // numrows
    dimensions[1] = width; // numcols
    if (png_get_color_type(png_ptr, info_ptr) & PNG_COLOR_MASK_ALPHA) {
        dimensions[2] = 4; // RGBA images
    } else if (png_get_color_type(png_ptr, info_ptr) & PNG_COLOR_MASK_COLOR) {
        dimensions[2] = 3; // RGB images
    } else {
        dimensions[2] = 1; // Greyscale images
    }

    if (float_result) {
        double max_value = (1 << bit_depth) - 1;

        numpy::array_view<float, 3> A(dimensions);

        for (png_uint_32 y = 0; y < height; y++) {
            png_byte *row = row_pointers[y];
            for (png_uint_32 x = 0; x < width; x++) {
                if (bit_depth == 16) {
                    png_uint_16 *ptr = &reinterpret_cast<png_uint_16 *>(row)[x * dimensions[2]];
                    for (png_uint_32 p = 0; p < (png_uint_32)dimensions[2]; p++) {
                        A(y, x, p) = (float)(ptr[p] / max_value);
                    }
                } else {
                    png_byte *ptr = &(row[x * dimensions[2]]);
                    for (png_uint_32 p = 0; p < (png_uint_32)dimensions[2]; p++) {
                        A(y, x, p) = (float)(ptr[p] / max_value);
                    }
                }
            }
        }

        result = A.pyobj();
    } else if (bit_depth == 16) {
        numpy::array_view<png_uint_16, 3> A(dimensions);

        for (png_uint_32 y = 0; y < height; y++) {
            png_byte *row = row_pointers[y];
            for (png_uint_32 x = 0; x < width; x++) {
                png_uint_16 *ptr = &reinterpret_cast<png_uint_16 *>(row)[x * dimensions[2]];
                for (png_uint_32 p = 0; p < (png_uint_32)dimensions[2]; p++) {
                    A(y, x, p) = ptr[p];
                }
            }
        }

        result = A.pyobj();
    } else if (bit_depth == 8) {
        numpy::array_view<png_byte, 3> A(dimensions);

        for (png_uint_32 y = 0; y < height; y++) {
            png_byte *row = row_pointers[y];
            for (png_uint_32 x = 0; x < width; x++) {
                png_byte *ptr = &(row[x * dimensions[2]]);
                for (png_uint_32 p = 0; p < (png_uint_32)dimensions[2]; p++) {
                    A(y, x, p) = ptr[p];
                }
            }
        }

        result = A.pyobj();
    } else {
        PyErr_SetString(PyExc_RuntimeError, "image has unknown bit depth");
        goto exit;
    }

    // free the png memory
    png_read_end(png_ptr, info_ptr);

    // For gray, return an x by y array, not an x by y by 1
    num_dims = (png_get_color_type(png_ptr, info_ptr) & PNG_COLOR_MASK_COLOR) ? 3 : 2;

    if (num_dims == 2) {
        PyArray_Dims dims = {dimensions, 2};
        PyObject *reshaped = PyArray_Newshape((PyArrayObject *)result, &dims, NPY_CORDER);
        Py_DECREF(result);
        result = reshaped;
    }

exit:
    if (png_ptr && info_ptr) {
#ifndef png_infopp_NULL
        png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
#else
        png_destroy_read_struct(&png_ptr, &info_ptr, png_infopp_NULL);
#endif
    }

    for (png_uint_32 row = 0; row < height; row++) {
        delete[] row_pointers[row];
    }

    if (PyErr_Occurred()) {
        Py_XDECREF(result);
        return NULL;
    } else {
        return result;
    }
}

const char *Py_read_png_float__doc__ =
    "read_png_float(file)\n"
    "\n"
    "Read in a PNG file, converting values to floating-point doubles\n"
    "in the range (0, 1)\n"
    "\n"
    "Parameters\n"
    "----------\n"
    "file : binary-mode file-like object\n";

static PyObject *Py_read_png_float(PyObject *self, PyObject *args, PyObject *kwds)
{
    return _read_png(args, true);
}

const char *Py_read_png_int__doc__ =
    "read_png_int(file)\n"
    "\n"
    "Read in a PNG file with original integer values.\n"
    "\n"
    "Parameters\n"
    "----------\n"
    "file : binary-mode file-like object\n";

static PyObject *Py_read_png_int(PyObject *self, PyObject *args, PyObject *kwds)
{
    return _read_png(args, false);
}

const char *Py_read_png__doc__ =
    "read_png(file)\n"
    "\n"
    "Read in a PNG file, converting values to floating-point doubles\n"
    "in the range (0, 1)\n"
    "\n"
    "Alias for read_png_float()\n"
    "\n"
    "Parameters\n"
    "----------\n"
    "file : binary-mode file-like object\n";

static PyMethodDef module_methods[] = {
    {"write_png", (PyCFunction)Py_write_png, METH_VARARGS|METH_KEYWORDS, Py_write_png__doc__},
    {"read_png", (PyCFunction)Py_read_png_float, METH_O, Py_read_png__doc__},
    {"read_png_float", (PyCFunction)Py_read_png_float, METH_O, Py_read_png_float__doc__},
    {"read_png_int", (PyCFunction)Py_read_png_int, METH_O, Py_read_png_int__doc__},
    {NULL}
};

extern "C" {

    static struct PyModuleDef moduledef = {
        PyModuleDef_HEAD_INIT,
        "_png",
        NULL,
        0,
        module_methods,
        NULL,
        NULL,
        NULL,
        NULL
    };

    PyMODINIT_FUNC PyInit__png(void)
    {
        PyObject *m;

        m = PyModule_Create(&moduledef);

        if (m == NULL) {
            return NULL;
        }

        import_array();

        if (PyModule_AddIntConstant(m, "PNG_FILTER_NONE", PNG_FILTER_NONE) ||
            PyModule_AddIntConstant(m, "PNG_FILTER_SUB", PNG_FILTER_SUB) ||
            PyModule_AddIntConstant(m, "PNG_FILTER_UP", PNG_FILTER_UP) ||
            PyModule_AddIntConstant(m, "PNG_FILTER_AVG", PNG_FILTER_AVG) ||
            PyModule_AddIntConstant(m, "PNG_FILTER_PAETH", PNG_FILTER_PAETH)) {
            return NULL;
        }


        return m;
    }
}
