#include <png.h>

// To remove a gcc warning
#ifdef _POSIX_C_SOURCE
#undef _POSIX_C_SOURCE
#endif

// TODO: Un CXX-ify this module
#include "CXX/Extensions.hxx"
#include "numpy/arrayobject.h"
#include "mplutils.h"

// the extension module
class _png_module : public Py::ExtensionModule<_png_module>
{
public:
    _png_module()
            : Py::ExtensionModule<_png_module>( "_png" )
    {
        add_varargs_method("write_png", &_png_module::write_png,
                           "write_png(buffer, width, height, fileobj, dpi=None)");
        add_varargs_method("read_png", &_png_module::read_png,
                           "read_png(fileobj)");
        initialize("Module to write PNG files");
    }

    virtual ~_png_module() {}

private:
    Py::Object write_png(const Py::Tuple& args);
    Py::Object read_png(const Py::Tuple& args);
};

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 *)"s#", 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);
}

// this code is heavily adapted from the paint license, which is in
// the file paint.license (BSD compatible) included in this
// distribution.  TODO, add license file to MANIFEST.in and CVS
Py::Object _png_module::write_png(const Py::Tuple& args)
{
  args.verify_length(4, 5);

  FILE *fp = NULL;
  bool close_file = false;
  Py::Object buffer_obj = Py::Object(args[0]);
  PyObject* buffer = buffer_obj.ptr();
  if (!PyObject_CheckReadBuffer(buffer)) {
    throw Py::TypeError("First argument must be an rgba buffer.");
  }

  const void* pixBufferPtr = NULL;
  Py_ssize_t pixBufferLength = 0;
  if (PyObject_AsReadBuffer(buffer, &pixBufferPtr, &pixBufferLength)) {
    throw Py::ValueError("Couldn't get data from read buffer.");
  }

  png_byte* pixBuffer = (png_byte*)pixBufferPtr;
  int width = (int)Py::Int(args[1]);
  int height = (int)Py::Int(args[2]);

  if (pixBufferLength < width * height * 4) {
    throw Py::ValueError("Buffer and width, height don't seem to match.");
  }

  Py::Object py_fileobj = Py::Object(args[3]);
  if (py_fileobj.isString()) {
    std::string fileName = Py::String(py_fileobj);
    const char *file_name = fileName.c_str();
    if ((fp = fopen(file_name, "wb")) == NULL)
      throw Py::RuntimeError( Printf("Could not open file %s", file_name).str() );
    close_file = true;
  } else if (PyFile_CheckExact(py_fileobj.ptr())) {
    fp = PyFile_AsFile(py_fileobj.ptr());
  }
  else {
    PyObject* write_method = PyObject_GetAttrString(py_fileobj.ptr(), "write");
    if (!(write_method && PyCallable_Check(write_method))) {
      Py_XDECREF(write_method);
      throw Py::TypeError("Object does not appear to be a 8-bit string path or a Python file-like object");
    }
    Py_XDECREF(write_method);
  }

  png_bytep *row_pointers = NULL;
  png_structp png_ptr = NULL;
  png_infop info_ptr = NULL;

  try {
    struct png_color_8_struct sig_bit;
    png_uint_32 row;

    row_pointers = new png_bytep[height];
    for (row = 0; row < (png_uint_32)height; ++row) {
      row_pointers[row] = pixBuffer + row * width * 4;
    }

    png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
    if (png_ptr == NULL) {
      throw Py::RuntimeError("Could not create write struct");
    }

    info_ptr = png_create_info_struct(png_ptr);
    if (info_ptr == NULL) {
      throw Py::RuntimeError("Could not create info struct");
    }

    if (setjmp(png_ptr->jmpbuf)) {
      throw Py::RuntimeError("Error building image");
    }

    if (fp) {
      png_init_io(png_ptr, fp);
    } else {
      png_set_write_fn(png_ptr, (void*)py_fileobj.ptr(),
		       &write_png_data, &flush_png_data);
    }
    png_set_IHDR(png_ptr, info_ptr,
		 width, height, 8,
		 PNG_COLOR_TYPE_RGB_ALPHA, PNG_INTERLACE_NONE,
		 PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    // Save the dpi of the image in the file
    if (args.size() == 5) {
      double dpi = Py::Float(args[4]);
      size_t dots_per_meter = (size_t)(dpi / (2.54 / 100.0));
      png_set_pHYs(png_ptr, info_ptr, dots_per_meter, dots_per_meter, PNG_RESOLUTION_METER);
    }

    // this a a color image!
    sig_bit.gray = 0;
    sig_bit.red = 8;
    sig_bit.green = 8;
    sig_bit.blue = 8;
    /* if the image has an alpha channel then */
    sig_bit.alpha = 8;
    png_set_sBIT(png_ptr, info_ptr, &sig_bit);

    png_write_info(png_ptr, info_ptr);
    png_write_image(png_ptr, row_pointers);
    png_write_end(png_ptr, info_ptr);
  } catch (...) {
      if (fp && close_file) fclose(fp);
      delete [] row_pointers;
      /* Changed calls to png_destroy_write_struct to follow
         http://www.libpng.org/pub/png/libpng-manual.txt.
         This ensures the info_ptr memory is released.
      */
      if (png_ptr && info_ptr) png_destroy_write_struct(&png_ptr, &info_ptr);
      throw;
  }

  png_destroy_write_struct(&png_ptr, &info_ptr);
  delete [] row_pointers;
  if (fp && close_file) fclose(fp);

  return Py::Object();
}


Py::Object
_png_module::read_png(const Py::Tuple& args) {

  args.verify_length(1);
  std::string fname = Py::String(args[0]);

  png_byte header[8];	// 8 is the maximum size that can be checked

  FILE *fp = fopen(fname.c_str(), "rb");
  if (!fp)
    throw Py::RuntimeError(Printf("_image_module::readpng could not open PNG file %s for reading", fname.c_str()).str());

  if (fread(header, 1, 8, fp) != 8)
    throw Py::RuntimeError("_image_module::readpng: error reading PNG header");
  if (png_sig_cmp(header, 0, 8))
    throw Py::RuntimeError("_image_module::readpng: file not recognized as a PNG file");


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

  if (!png_ptr)
    throw Py::RuntimeError("_image_module::readpng:  png_create_read_struct failed");

  png_infop info_ptr = png_create_info_struct(png_ptr);
  if (!info_ptr)
    throw Py::RuntimeError("_image_module::readpng:  png_create_info_struct failed");

  if (setjmp(png_jmpbuf(png_ptr)))
    throw Py::RuntimeError("_image_module::readpng:  error during init_io");

  png_init_io(png_ptr, fp);
  png_set_sig_bytes(png_ptr, 8);

  png_read_info(png_ptr, info_ptr);

  png_uint_32 width = info_ptr->width;
  png_uint_32 height = info_ptr->height;
  bool do_gray_conversion = (info_ptr->bit_depth < 8 &&
                             info_ptr->color_type == PNG_COLOR_TYPE_GRAY);

  int bit_depth = info_ptr->bit_depth;
  if (bit_depth == 16) {
    png_set_strip_16(png_ptr);
  } else if (bit_depth < 8) {
    png_set_packing(png_ptr);
  }

  // convert misc color types to rgb for simplicity
  if (info_ptr->color_type == PNG_COLOR_TYPE_GRAY ||
      info_ptr->color_type == PNG_COLOR_TYPE_GRAY_ALPHA) {
    png_set_gray_to_rgb(png_ptr);
  } else if (info_ptr->color_type == PNG_COLOR_TYPE_PALETTE) {
    png_set_palette_to_rgb(png_ptr);
  }

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

  bool rgba = info_ptr->color_type == PNG_COLOR_TYPE_RGBA;
  if ( (info_ptr->color_type != PNG_COLOR_TYPE_RGB) && !rgba) {
    std::cerr << "Found color type " << (int)info_ptr->color_type  << std::endl;
    throw Py::RuntimeError("_image_module::readpng: cannot handle color_type");
  }

  /* read file */
  if (setjmp(png_jmpbuf(png_ptr)))
    throw Py::RuntimeError("_image_module::readpng: error during read_image");

  png_bytep *row_pointers = new png_bytep[height];
  png_uint_32 row;

  for (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);

  npy_intp dimensions[3];
  dimensions[0] = height;  //numrows
  dimensions[1] = width;   //numcols
  dimensions[2] = 4;

  PyArrayObject *A = (PyArrayObject *) PyArray_SimpleNew(3, dimensions, PyArray_FLOAT);

  if (do_gray_conversion) {
    float max_value = (float)((1L << bit_depth) - 1);
    for (png_uint_32 y = 0; y < height; y++) {
      png_byte* row = row_pointers[y];
      for (png_uint_32 x = 0; x < width; x++) {
        float value = row[x] / max_value;
        size_t offset = y*A->strides[0] + x*A->strides[1];
        *(float*)(A->data + offset + 0*A->strides[2]) = value;
        *(float*)(A->data + offset + 1*A->strides[2]) = value;
        *(float*)(A->data + offset + 2*A->strides[2]) = value;
        *(float*)(A->data + offset + 3*A->strides[2]) = 1.0f;
      }
    }
  } else {
    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 = (rgba) ? &(row[x*4]) : &(row[x*3]);
        size_t offset = y*A->strides[0] + x*A->strides[1];
        *(float*)(A->data + offset + 0*A->strides[2]) = (float)(ptr[0]/255.0);
        *(float*)(A->data + offset + 1*A->strides[2]) = (float)(ptr[1]/255.0);
        *(float*)(A->data + offset + 2*A->strides[2]) = (float)(ptr[2]/255.0);
        *(float*)(A->data + offset + 3*A->strides[2]) = rgba ? (float)(ptr[3]/255.0) : 1.0f;
      }
    }
  }

  //free the png memory
  png_read_end(png_ptr, info_ptr);
  png_destroy_read_struct(&png_ptr, &info_ptr, png_infopp_NULL);
  fclose(fp);
  for (row = 0; row < height; row++)
    delete [] row_pointers[row];
  delete [] row_pointers;
  return Py::asObject((PyObject*)A);
}

extern "C"
    DL_EXPORT(void)
    init_png(void)
{
    import_array();

    static _png_module* _png = NULL;
    _png = new _png_module;
}