/*
 * Copyright 2022 Rive
 */

#include "write_png_file.hpp"

#include "png.h"
#include "zlib.h"
#include <errno.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <vector>

void WritePNGFile(
    uint8_t* pixels,
    int width,
    int height,
    bool flipY,
    const char* file_name,
    PNGCompression compression)
{
  FILE* fp = fopen(file_name, "wb");
  if (!fp)
  {
    fprintf(
        stderr,
        "WritePNGFile: File %s could not be opened for writing "
        "(errno=%i)\n",
        file_name,
        errno);
    abort();
  }

  auto png = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);

  if (!png)
  {
    fprintf(stderr, "WritePNGFile: png_create_write_struct failed\n");
    abort();
  }

  if (compression == PNGCompression::fast_rle)
  {
    // RLE with SUB gets best performance with our content.
    png_set_compression_strategy(png, Z_RLE);
    png_set_filter(png, 0, PNG_FILTER_SUB);
  }

  auto info = png_create_info_struct(png);
  if (!info)
  {
    fprintf(stderr, "WritePNGFile: png_create_info_struct failed\n");
    abort();
  }

  if (setjmp(png_jmpbuf(png)))
  {
    fprintf(stderr, "WritePNGFile: Error during init_io\n");
    abort();
  }

  png_init_io(png, fp);

  // Write header.
  if (setjmp(png_jmpbuf(png)))
  {
    fprintf(stderr, "WritePNGFile: Error during writing header\n");
    abort();
  }

  png_set_IHDR(
      png,
      info,
      width,
      height,
      8,
      PNG_COLOR_TYPE_RGB_ALPHA,
      PNG_INTERLACE_NONE,
      PNG_COMPRESSION_TYPE_BASE,
      PNG_FILTER_TYPE_BASE);

  png_write_info(png, info);

  // Write bytes.
  if (setjmp(png_jmpbuf(png)))
  {
    fprintf(stderr, "WritePNGFile: Error during writing bytes\n");
    abort();
  }

  std::vector<uint8_t*> rows(height);
  for (int y = 0; y < height; ++y)
  {
    rows[y] = pixels + (flipY ? height - 1 - y : y) * width * 4;
  }
  png_write_image(png, rows.data());

  // End write.
  if (setjmp(png_jmpbuf(png)))
  {
    fprintf(stderr, "WritePNGFile: Error during end of write");
    abort();
  }

  png_write_end(png, NULL);
  png_destroy_write_struct(&png, &info);

  fclose(fp);
}

static void append_png_data_chunk(
    png_structp png,
    png_bytep data,
    png_size_t length)
{
  auto encodedData = reinterpret_cast<std::vector<uint8_t>*>(
      png_get_io_ptr(png));
  encodedData->insert(encodedData->end(), data, data + length);
}

static void flush_png_data(png_structp png) {}

std::vector<uint8_t> EncodePNGToBuffer(
    uint32_t width,
    uint32_t height,
    uint8_t* imageDataRGBA,
    PNGCompression compression)
{
  // Send the image name.
  auto png = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);

  if (!png)
  {
    fprintf(stderr, "WritePNGFile: png_create_write_struct failed\n");
    abort();
  }

  if (compression == PNGCompression::fast_rle)
  {
    // RLE with SUB gets best performance with our content.
    png_set_compression_level(png, 6);
    png_set_compression_strategy(png, Z_RLE);
    png_set_compression_strategy(png, Z_RLE);
    png_set_filter(png, 0, PNG_FILTER_SUB);
  }

  auto info = png_create_info_struct(png);
  if (!info)
  {
    fprintf(stderr, "WritePNGFile: png_create_info_struct failed\n");
    abort();
  }

  if (setjmp(png_jmpbuf(png)))
  {
    fprintf(stderr, "WritePNGFile: Error during init_io\n");
    abort();
  }

  std::vector<uint8_t> encodedData;
  png_set_write_fn(png, &encodedData, &append_png_data_chunk, &flush_png_data);

  // Write header.
  if (setjmp(png_jmpbuf(png)))
  {
    fprintf(stderr, "WritePNGFile: Error during writing header\n");
    abort();
  }

  png_set_IHDR(
      png,
      info,
      width,
      height,
      8,
      PNG_COLOR_TYPE_RGB_ALPHA,
      PNG_INTERLACE_NONE,
      PNG_COMPRESSION_TYPE_BASE,
      PNG_FILTER_TYPE_BASE);

  png_write_info(png, info);

  // Write bytes.
  if (setjmp(png_jmpbuf(png)))
  {
    fprintf(stderr, "WritePNGFile: Error during writing bytes\n");
    abort();
  }

  std::vector<uint8_t*> rows(height);
  for (uint32_t y = 0; y < height; ++y)
  {
    rows[y] = imageDataRGBA + (height - 1 - y) * width * 4;
  }
  png_write_image(png, rows.data());

  // End write.
  if (setjmp(png_jmpbuf(png)))
  {
    fprintf(stderr, "WritePNGFile: Error during end of write");
    abort();
  }

  png_write_end(png, NULL);
  png_destroy_write_struct(&png, &info);

  return encodedData;
}