/* -*- mia-c++  -*-
 *
 * This file is part of MIA - a toolbox for medical image analysis
 * Copyright (c) Leipzig, Madrid 1999-2017 Gert Wollny
 *
 * MIA is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with MIA; if not, see <http://www.gnu.org/licenses/>.
 *
 */

#include <png.h>
#include <sstream>
#include <cassert>
#include <boost/shared_array.hpp>
#include <mia/core/file.hh>
#include <mia/core/filter.hh>
#include <mia/core/msgstream.hh>
//#include <mia/core/errormacro.hh>
#include <mia/2d/rgbimageio.hh>



NS_BEGIN(RGBIMAGEIO_2D_PNG)


using namespace mia;
using namespace std;
using namespace boost;

class CPNG2DImageIOPlugin : public C2DRGBImageIOPlugin
{
public:
       CPNG2DImageIOPlugin();
private:
       PData do_load(const string& fname) const;
       bool do_save(const string& fname, const Data& data) const;
       const string do_get_descr() const;
};

CPNG2DImageIOPlugin::CPNG2DImageIOPlugin():
       C2DRGBImageIOPlugin("png")
{
       add_supported_type(it_ubyte);
       add_property(io_plugin_property_has_attributes);
       add_suffix(".png");
       add_suffix(".PNG");
}

CPNG2DImageIOPlugin::PData  CPNG2DImageIOPlugin::do_load(const string& fname) const
{
       CInputFile f(fname);

       if (!f) {
              cvdebug() << fname << " not found\n";
              return CPNG2DImageIOPlugin::PData();
       }

       const size_t sig_number = 8;
       png_byte header[sig_number];

       if (fread(header, 1, sig_number, f) != sig_number) {
              return CPNG2DImageIOPlugin::PData();
       }

       if (png_sig_cmp(header, 0, sig_number)) {
              cvdebug() << "png: signature not found\n";
              return CPNG2DImageIOPlugin::PData();
       }

       png_structp png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);

       if (!png_ptr)
              throw runtime_error("unable to allocate png read_struct");

       png_infop info_ptr = png_create_info_struct(png_ptr);

       if (!info_ptr) {
              png_destroy_read_struct(&png_ptr, (png_infopp)NULL, (png_infopp)NULL);
              throw runtime_error("unable to allocate png info ptr");
       }

       png_infop end_info = png_create_info_struct(png_ptr);

       if (!end_info) {
              png_destroy_read_struct(&png_ptr, &info_ptr, (png_infopp)NULL);
              throw runtime_error("unable to allocate png end info ptr");
       }

       if (setjmp(png_jmpbuf(png_ptr))) {
              png_destroy_read_struct(&png_ptr, &info_ptr, &end_info);
              throw runtime_error("Error reading png file");
       }

       png_init_io(png_ptr, f);
       png_set_sig_bytes(png_ptr, sig_number);
       png_set_packing(png_ptr);
       png_read_info(png_ptr, info_ptr);
       png_uint_32 width;
       png_uint_32 height;
       int bit_depth;
       int color_type;
       int interlace_type;
       png_get_IHDR(png_ptr, info_ptr, &width, &height, &bit_depth, &color_type,
                    &interlace_type, NULL, NULL);

       if (color_type != PNG_COLOR_TYPE_RGB) {
              png_destroy_read_struct(&png_ptr, &info_ptr, &end_info);
              stringstream errmsg;
              errmsg << "CPNG2DImageIO::load: image type is '";

              switch (color_type) {
              case PNG_COLOR_TYPE_GRAY:
                     errmsg << "GRAY";
                     break;

              case PNG_COLOR_TYPE_RGB_ALPHA:
                     errmsg << "RGBA";
                     break;

              case PNG_COLOR_TYPE_GRAY_ALPHA:
                     errmsg << "GRAY+ALPHA";
                     break;

              case PNG_COLOR_TYPE_PALETTE:
                     errmsg << "COLOR+PALETTE";
                     break;

              default:
                     assert(!"RGB and PALETTE should not be rejected");
              }

              errmsg <<  "' but only support RGB images are supported.";
              throw invalid_argument(errmsg.str());
       }

       if (bit_depth != 8 )
              throw create_exception<invalid_argument>("PNG-RGB: Only 8-bit per color components are supported, but got ",  bit_depth, " bit");

       png_read_update_info(png_ptr, info_ptr);
       shared_array<png_bytep> row_pointers(new png_bytep[height]);
       size_t row_length = png_get_rowbytes(png_ptr, info_ptr);
       cvdebug() << "png: image = " << width << "x" << height << "@" << bit_depth << "\n";
       cvdebug() << "png:row_length = " << row_length << "\n";

       for (png_uint_32 row = 0; row < height; row++) {
              row_pointers[row] = static_cast<png_bytep>(png_malloc(png_ptr, row_length));
       }

       png_read_image(png_ptr, row_pointers.get());
       png_read_end(png_ptr, info_ptr);
       png_destroy_read_struct(&png_ptr, &info_ptr, &end_info);
       C2DBounds size(width, height);
       CRGB2DImage *result = new CRGB2DImage(size);
       auto p = result->pixel();

       for (png_uint_32 row = 0; row < height; row++, p += row_length) {
              copy(row_pointers[row], row_pointers[row] + row_length, p);
              free(row_pointers[row]);
       }

       return CPNG2DImageIOPlugin::PData(result);
}





bool CPNG2DImageIOPlugin::do_save(const string& fname, const CRGB2DImage& image) const
{
       TRACE_FUNCTION;
       COutputFile f(fname);

       if (!f) {
              cverr() << "CPNG2DImageIO::save:unable to open output file:" << fname << "\n";
              return false;
       }

       png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING,
                             NULL, NULL, NULL);

       if (!png_ptr)
              throw runtime_error("unable to create png write struct");

       png_infop info_ptr = png_create_info_struct(png_ptr);

       if (!info_ptr) {
              png_destroy_write_struct(&png_ptr,  NULL);
              throw runtime_error("unable to create png info struct");
       }

       if (setjmp(png_jmpbuf(png_ptr))) {
              /* If we get here, we had a problem reading the file */
              png_destroy_write_struct(&png_ptr, &info_ptr);
              throw runtime_error("Error writing png file");
       }

       png_init_io(png_ptr, f);
       png_set_IHDR(png_ptr, info_ptr, image.get_size().x, image.get_size().y,
                    8, PNG_COLOR_TYPE_RGB,
                    PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
       png_write_info(png_ptr, info_ptr);
       png_set_packing(png_ptr);
       int line_stride = image.get_size().x * 3;
       auto p = image.pixel();
       vector<unsigned char> buffer(line_stride);

       for (auto i = 0u; i < image.get_size().y; i++, p += line_stride) {
              copy(p, p + line_stride, buffer.begin());
              png_write_row(png_ptr, &buffer[0]);
       }

       png_write_end(png_ptr, info_ptr);
       png_destroy_write_struct(&png_ptr, &info_ptr);
       return true;
}

const string CPNG2DImageIOPlugin::do_get_descr() const
{
       return "a 2d RGB image io plugin for png images";
}

extern "C" EXPORT  CPluginBase *get_plugin_interface()
{
       return new CPNG2DImageIOPlugin();
}

NS_END
