#include <psemek/gfx/pixmap.hpp>

#include <psemek/log/log.hpp>
#include <psemek/util/at_scope_exit.hpp>
#include <psemek/util/exception.hpp>

#include <png.h>

namespace psemek::gfx
{

	namespace
	{

		template <typename Pixel>
		basic_pixmap<Pixel> read_png_impl(io::istream & is, bool monochrome)
		{
			png_error_ptr error = [](png_structp, png_const_charp str)
			{
				throw util::exception(str);
			};

			png_error_ptr warn = [](png_structp, png_const_charp str)
			{
				log::warning() << str;
			};

			auto png = png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, error, warn);
			if (!png) throw util::exception("png_create_read_struct returned null");

			png_set_error_fn(png, nullptr, error, warn);

			png_infop info = nullptr;
			png_infop end_info = nullptr;

			[[maybe_unused]] auto png_dtor = util::at_scope_exit([&]{
				png_destroy_read_struct(&png, &info, &end_info);
			});

			info = png_create_info_struct(png);
			if (!info) throw util::exception("png_create_info_struct returned null");

			end_info = png_create_info_struct(png);
			if (!end_info) throw util::exception("png_create_info_struct returned null");

			png_rw_ptr read = [](png_structp png, png_bytep ptr, size_t count){
				reinterpret_cast<io::istream *>(png_get_io_ptr(png))->read(reinterpret_cast<char *>(ptr), count);
			};
			png_set_read_fn(png, &is, read);

			png_read_info(png, info);

			auto const width = png_get_image_width(png, info);
			auto const height = png_get_image_height(png, info);
			auto const color_type = png_get_color_type(png, info);
			auto const bit_depth = png_get_bit_depth(png, info);

			if (bit_depth == 16)
				png_set_strip_16(png);

			if (monochrome && color_type != PNG_COLOR_TYPE_GRAY)
				throw util::exception("Invalid color type for monochrome PNG");

			if (color_type == PNG_COLOR_TYPE_PALETTE)
				png_set_palette_to_rgb(png);

			if (color_type == PNG_COLOR_TYPE_GRAY && bit_depth < 8)
				png_set_expand_gray_1_2_4_to_8(png);

			if (color_type == PNG_COLOR_TYPE_RGB)
				png_set_add_alpha(png, 0xff, PNG_FILLER_AFTER);

			png_read_update_info(png, info);

			basic_pixmap<Pixel> result({width, height});

			auto const row_bytes = png_get_rowbytes(png, info);
			if (row_bytes != width * sizeof(Pixel))
				throw util::exception("PNG row bytes mismatch");

			for (std::uint32_t i = 0; i < height; ++i)
				png_read_row(png, reinterpret_cast<png_bytep>(result.data() + i * width), nullptr);

			return result;
		}

		void write_png_impl(pixmap_rgba const & p, io::ostream & os)
		{
			png_error_ptr error = [](png_structp, png_const_charp str)
			{
				throw util::exception(str);
			};

			png_error_ptr warn = [](png_structp, png_const_charp str)
			{
				log::warning() << str;
			};

			auto png = png_create_write_struct(PNG_LIBPNG_VER_STRING, nullptr, error, warn);
			if (!png) throw util::exception("png_create_write_struct returned null");

			png_set_error_fn(png, nullptr, error, warn);

			png_infop info = nullptr;

			[[maybe_unused]] auto png_dtor = util::at_scope_exit([&]{
				png_destroy_write_struct(&png, &info);
			});

			info = png_create_info_struct(png);
			if (!info) throw util::exception("png_create_info_struct returned null");

			png_rw_ptr write = [](png_structp png, png_bytep ptr, size_t count){
				reinterpret_cast<io::ostream *>(png_get_io_ptr(png))->write(reinterpret_cast<char const *>(ptr), count);
			};
			png_flush_ptr flush = [](png_structp png){
				reinterpret_cast<io::ostream *>(png_get_io_ptr(png))->flush();
			};
			png_set_write_fn(png, &os, write, flush);

			png_set_IHDR(png, info, p.width(), p.height(), 8, PNG_COLOR_TYPE_RGBA, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);

			png_write_info(png, info);

			for (std::uint32_t i = 0; i < p.height(); ++i)
				png_write_row(png, reinterpret_cast<png_byte const *>(p.data() + p.width() * i));
		}

	}

	pixmap_rgba read_png(io::istream && is)
	{
		return read_png_impl<color_rgba>(is, false);
	}

	pixmap_monochrome read_png_monochrome(io::istream && is)
	{
		return read_png_impl<std::uint8_t>(is, true);
	}

	void write_png(pixmap_rgba const & p, io::ostream && os)
	{
		write_png_impl(p, os);
	}

}