/*
 * Copyright 2023 Rive
 */

#include "rive/decoders/bitmap_decoder.hpp"
#include "png.h"
#include <algorithm>
#include <cassert>
#include <string.h>
#include <setjmp.h>

struct EncodedImageBuffer
{
  const uint8_t* bytes;
  size_t position;
  size_t size;
};

static void ReadDataFromMemory(
    png_structp png_ptr,
    png_bytep outBytes,
    png_size_t byteCountToRead)
{
  png_voidp a = png_get_io_ptr(png_ptr);
  if (a == nullptr)
  {
    return;
  }
  EncodedImageBuffer& stream = *(EncodedImageBuffer*)a;

  size_t bytesRead = std::min(byteCountToRead, (stream.size - stream.position));
  memcpy(outBytes, stream.bytes + stream.position, bytesRead);
  stream.position += bytesRead;

  if ((png_size_t)bytesRead != byteCountToRead)
  {
    // Report image error?
  }
}

std::unique_ptr<Bitmap> DecodePng(const uint8_t bytes[], size_t byteCount)
{
  png_structp png_ptr;
  png_infop info_ptr;
  png_uint_32 width, height;
  int bit_depth, color_type, interlace_type;
  std::unique_ptr<uint8_t[]> pixelBuffer;
  std::unique_ptr<png_bytep[]> rowsPointer;

  png_ptr = png_create_read_struct(
      PNG_LIBPNG_VER_STRING,
      nullptr,
      nullptr,
      nullptr);

  if (png_ptr == nullptr)
  {
    printf("DecodePng - libpng failed (png_create_read_struct).");
    return nullptr;
  }

  info_ptr = png_create_info_struct(png_ptr);
  if (info_ptr == NULL)
  {
    png_destroy_read_struct(&png_ptr, (png_infopp)NULL, (png_infopp)NULL);
    printf("DecodePng - libpng failed (png_create_info_struct).");
    return nullptr;
  }

  if (setjmp(png_jmpbuf(png_ptr)))
  {
    png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
    return nullptr;
  }

  EncodedImageBuffer stream = {bytes, 0, byteCount};

  png_set_read_fn(png_ptr, &stream, ReadDataFromMemory);

  png_read_info(png_ptr, info_ptr);

  png_get_IHDR(
      png_ptr,
      info_ptr,
      &width,
      &height,
      &bit_depth,
      &color_type,
      &interlace_type,
      NULL,
      NULL);

  if (color_type == PNG_COLOR_TYPE_PALETTE ||
      (color_type == PNG_COLOR_TYPE_GRAY && bit_depth < 8))
  {
    // expand to 3 or 4 channels
    png_set_expand(png_ptr);
  }

  png_bytep trns = 0;
  int trnsCount = 0;
  if (png_get_valid(png_ptr, info_ptr, PNG_INFO_tRNS))
  {
    png_get_tRNS(png_ptr, info_ptr, &trns, &trnsCount, 0);
    png_set_expand(png_ptr);
  }

  if (bit_depth == 16)
  {
    png_set_strip_16(png_ptr);
  }

  if (color_type == PNG_COLOR_TYPE_GRAY ||
      color_type == PNG_COLOR_TYPE_GRAY_ALPHA)
  {
    png_set_gray_to_rgb(png_ptr);
  }

  png_read_update_info(png_ptr, info_ptr);
  uint8_t channels = png_get_channels(png_ptr, info_ptr);

  size_t pixelBufferSize = static_cast<size_t>(width) *
                           static_cast<size_t>(height) *
                           static_cast<size_t>(channels);
  pixelBuffer = std::make_unique<uint8_t[]>(pixelBufferSize);
  const uint8_t* pixelBufferEnd = pixelBuffer.get() + pixelBufferSize;

  rowsPointer = std::make_unique<png_bytep[]>(height);
  png_bytep* rows = rowsPointer.get();
  size_t rowStride = (size_t)width * (size_t)channels;
  uint8_t* rowWrite = pixelBuffer.get();
  for (png_uint_32 row = 0; row < height; row++)
  {
    rows[row] = rowWrite;
    rowWrite += rowStride;
    if (rowWrite > pixelBufferEnd)
    {
      // Read would overflow.
      return nullptr;
    }
  }
  png_read_image(png_ptr, rows);
  png_read_end(png_ptr, info_ptr);

  png_destroy_read_struct(&png_ptr, &info_ptr, (png_infopp) nullptr);

  Bitmap::PixelFormat pixelFormat;
  assert(channels == 3 || channels == 4);
  switch (channels)
  {
    case 4:
      pixelFormat = Bitmap::PixelFormat::RGBA;
      break;
    case 3:
      pixelFormat = Bitmap::PixelFormat::RGB;
      break;
  }

  assert(pixelBufferSize == height * width * channels);
  return std::make_unique<Bitmap>(
      width,
      height,
      pixelBufferSize,
      pixelFormat,
      std::move(pixelBuffer));
}