/* See {jsjpeg_image.h} */
/* Last edited on 2016-12-13 17:23:09 by stolfilocal */

/* Created by R. Minetto (IC-UNICAMP) as {ipng.c} sometime in 2008--2009. */
/* Adapted by J. Stolfi (IC-UNICMP) on 2011-05-14. */

#define _GNU_SOURCE
#include <stdlib.h>
#include <math.h>
#include <string.h>
#include <stdio.h>
#include <assert.h>
#include <png.h>

#include <bool.h>
#include <affirm.h>
#include <jsfile.h>
#include <jspnm.h>
#include <jspnm_image.h>
#include <jspng_image.h>

/* INTERNAL PROTOTYPES */

void jspng_dump_info(FILE *wr, const char *func, char *label, png_structp pr, png_infop pi);

/* IMPLEMENTATIONS */

pnm_image_t *jspng_image_read (char *name, double *gammaP, bool_t verbose)
  {
    FILE *rd = open_read (name, verbose);
    pnm_image_t *img = jspng_image_fread(rd, gammaP, verbose);
    if (rd != stdin) { fclose(rd); }
    return img;
  }

void jspng_image_write (char *name, pnm_image_t *img, double gamma, bool_t verbose)
  {
    FILE *wr = open_write (name, verbose);
    jspng_image_fwrite(wr, img, gamma, verbose);
    if ((wr != stdout) && (wr != stderr)) { fclose(wr); }
  }

#define PNG_MAGIC_BYTES 8
/* Number of bytes in the PNG file signature. */

#define PNG_MAX_SIZE (2147483647u)
/* Max width and height of a PNG file that can be read by this interface. */

#define PNG_MAX_CHNS (4)
/* Max channels in a PNG image that can be read with this interface. */
   

pnm_image_t *jspng_image_fread (FILE *rd, double *gammaP, bool_t verbose)
  {
    bool_t debug = FALSE;
    
    /* Check the PNG signature: */
    unsigned char sig[PNG_MAGIC_BYTES];
    fread(sig, 1, PNG_MAGIC_BYTES, rd);
    demand(png_sig_cmp(sig, 0, PNG_MAGIC_BYTES) == 0, "not a PNG file");

    /* Allocate the reader structures: */
    png_structp pr = notnull(png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL), "no mem");
    png_infop pi = notnull(png_create_info_struct(pr), "no mem");

    /* Attach the file handle to the PNG reader, tell it we already read the sig: */
    png_init_io(pr, rd);
    png_set_sig_bytes(pr, PNG_MAGIC_BYTES);

    /* Read the header: */
    png_read_info(pr, pi);
    if (debug) { jspng_dump_info(stderr, __FUNCTION__, "before transforms", pr, pi); }

    /* Get main attributes: */
    png_uint_32 chns = png_get_channels(pr, pi);
    demand((chns > 0) && (chns <= PNG_MAX_CHNS), "too many channels"); 
    
    png_uint_32 cols = png_get_image_width(pr, pi);
    demand(cols <= PNG_MAX_SIZE, "too many cols"); 
    
    png_uint_32 rows = png_get_image_height(pr, pi);
    demand(rows <= PNG_MAX_SIZE, "too many rows"); 
    
    /* Get the color space and encoding: */
    uint8_t color_type = png_get_color_type(pr, pi);
    
    /* Get the number of bits per sample {bits} in the file: */
    uint8_t bits = png_get_bit_depth(pr, pi); /* Bits per sample in file. */
    demand((bits == 1) || (bits == 2) || (bits == 4) || (bits == 8) || (bits == 16), "invalid bits per sample in file"); 
    if (color_type == PNG_COLOR_TYPE_GRAY)
      { /* Any bit size is OK. */ }
    else if (color_type == PNG_COLOR_TYPE_GRAY_ALPHA)
      { demand(bits >= 8, "invalid bits per sample for gray+alpha color type"); }
    else if (color_type ==  PNG_COLOR_TYPE_PALETTE)
      { demand(bits <= 8, "invalid bits per sample for palette color type"); }
    else if (color_type == PNG_COLOR_TYPE_RGB)
      { demand(bits >= 8, "invalid bits per sample for RGB color type"); }
    else if (color_type == PNG_COLOR_TYPE_RGB_ALPHA)
      { demand(bits >= 8, "invalid bits per sample for RGB+alpha color type"); }
    else
      { demand(FALSE, "invalid color type"); }
    
    /* Find number {buse} of significant bits in each sample: */
    uint8_t buse;
    { 
      /* Get the number of significant bits in each file sample, per channel: */
      png_color_8p sBIT; 
      if (! png_get_sBIT(pr, pi, &sBIT))
        { /* Assume that all bits are significant: */
          buse = bits;
        }
      else
        { /* Check the {sBIT} record: */
          demand((sBIT->red   > 0) && (sBIT->red   <= bits), "inconsistent significant bits - red");
          demand((sBIT->green > 0) && (sBIT->green <= bits), "inconsistent significant bits - green");
          demand((sBIT->blue  > 0) && (sBIT->blue  <= bits), "inconsistent significant bits - blue");
          demand((sBIT->gray  > 0) && (sBIT->gray  <= bits), "inconsistent significant bits - gray");
          demand((sBIT->alpha > 0) && (sBIT->alpha <= bits), "inconsistent significant bits - alpha");
          /* Find max number of significant bits in any channel: */
          buse = 0;
          if (sBIT->red   > buse) { buse = sBIT->red;   }
          if (sBIT->green > buse) { buse = sBIT->green; }
          if (sBIT->blue  > buse) { buse = sBIT->blue;  }
          if (sBIT->gray  > buse) { buse = sBIT->gray;  }
          if (sBIT->alpha > buse) { buse = sBIT->alpha; }
        }
    }
    assert((buse > 0) && (buse <= bits));

    /* Get the gamma exponent of the file (but do not do gamma correction): */
    if (! png_get_gAMA(pr, pi, gammaP)) { (*gammaP) = NAN; }

    /* -- INPUT TRANSFORMATION REQUESTS -------------------------------- */

    /* Request expansion of 1,2,4-bit samples to one byte: */
    /* This will leave each sample in the LOWER {bits} bits of each byte. */
    if (bits < 8) { png_set_packing(pr); } 

    /* Request conversion of palette colors to RGB or RGBA: */
    if (color_type == PNG_COLOR_TYPE_PALETTE) { png_set_palette_to_rgb(pr); }

    /* If palette with a transparent color, request conversion to RGBA: */
    if (png_get_valid(pr, pi, PNG_INFO_tRNS)) { png_set_tRNS_to_alpha(pr); }
    
    /* Request de-interlacing: */
    (void)png_set_interlace_handling(pr);

    /* ------------------------------------------------------------------ */

    /* Update the derived info fields: */
    png_read_update_info(pr, pi);
    if (debug) { jspng_dump_info(stderr, __FUNCTION__, "after transforms", pr, pi); }

    /* Allocate the image in memory: */
    pnm_image_t *img = pnm_image_new((int)cols, (int)rows, (int)chns);
    img->maxval = (pnm_sample_t)((1 << buse) - 1);

    /* Samples per row in PNG and PNM images: */
    int spr = (int)(cols*chns); /* Samples per row. */

    /* Allocate scanlines of PNG image: */
    unsigned int png_row_bytes = (unsigned int)png_get_rowbytes(pr, pi);
    png_bytep *png_arrayP = notnull(malloc(rows*sizeof(png_bytep)), "no mem");
    int row;
    for (row = 0; row < rows; row++) { png_arrayP[row] = notnull(malloc(png_row_bytes), "no mem"); }

    /* Read whole image: */
    png_read_image(pr, png_arrayP);

    uint8_t shift = (uint8_t)(bits - buse);
    if (debug) { fprintf(stderr, "shift = %u\n", (unsigned int)shift); } 
    for (row = 0; row < rows; row++)
      { 
        png_bytep png_P = png_arrayP[row];
        pnm_sample_t *pnm_P = img->smp[row];
        int k;
        for (k = 0; k < spr; k++)
          { int png_smp = (*png_P); png_P++;
            if (bits > 8)
              { /* PNG file has two bytes per sample: */
                png_smp = (png_smp << 8) | (*png_P); png_P++;
              }
            int smp = (png_smp >> shift);
            if ((smp << shift) != png_smp)
              { fprintf
                  ( stderr, "png_smp = %u  smp = %u  smp shifted = %u\n", 
                    (unsigned int)png_smp, (unsigned int)smp, (unsigned int)(smp << shift)
                    );
              }
            assert((smp << shift) == png_smp);
            assert(smp <= img->maxval);
            (*pnm_P) = (pnm_sample_t)smp;
            pnm_P++;
          }
      }    
    for (row = 0; row < rows; row++) { free(png_arrayP[row]); png_arrayP[row] = NULL; }
    free(png_arrayP);
    /* png_destroy_read_struct(&pr, &pi, NULL); */
    png_free_data(pr, pi, PNG_FREE_ALL, -1);
    return img; 
  }
  
void jspng_image_fwrite (FILE *wr, pnm_image_t *img, double gamma, bool_t verbose)
  {
    bool_t debug = FALSE;
    
    png_uint_32 chns = img->chns;
    png_uint_32 cols = img->cols;
    png_uint_32 rows = img->rows;
    
    pnm_sample_t imaxval = img->maxval;  /* Image's {maxval}. */
    demand(imaxval + 0 <= 65535, "maxval is too big");
    demand(imaxval > 0, "maxval is zero");

    /* Create {libpng} work records {ppng, pinf}. */
    /* Avoid {setjmp/longjmp} complexity for now, use default error handling. */
    png_struct *ppng = notnull(png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL), "no mem");
    png_info *pinf = notnull(png_create_info_struct(ppng), "no mem");

    /* Assign {wr} as the output file for writes to {ppng}: */
    png_init_io(ppng, wr);

    /* Set the {zlib} compression level: */
    png_set_compression_level(ppng, 7); /* Level 6 is default, level 9 is maximum. */

     /* Select the file bit depth: */
    png_byte bits = (imaxval == 1 ? 1 : (imaxval <= 3 ? 2 : (imaxval <= 15 ? 4: (imaxval <= 255 ? 8 : 16))));
    
    /* Select a scaling factor so that {maxval} always goes to {2^bits-1}. */
    int omaxval = (1 << bits) - 1;
    double scale = ((double)omaxval)/((double)imaxval);
    if (verbose) 
      { if (imaxval != omaxval) { fprintf(stderr, "rescaling from [0..%d] to [0..%d]\n", imaxval, omaxval); }
        if (scale != floor(scale)) { fprintf(stderr, "there may be roundoff errors\n"); }
      }
    
    /* Select output image type: */
    png_byte ctype;
    if (chns == 1)
      { ctype = PNG_COLOR_TYPE_GRAY; }
    else if (chns == 2)
      { ctype = PNG_COLOR_TYPE_GRAY_ALPHA; }
    else if (chns == 3)
      { ctype = PNG_COLOR_TYPE_RGB; }
    else if (chns == 4)
      { ctype = PNG_COLOR_TYPE_RGB_ALPHA; }
    else
      { demand(FALSE, "invalid number of channels"); }

    png_set_IHDR(ppng, pinf, cols, rows, bits, ctype, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
    if (debug) { jspng_dump_info(stderr, __FUNCTION__, "before write_info", ppng, pinf); }

    /* Provide a "gAMA" chunk. */
    if (! isnan(gamma) && (gamma > 0)) { png_set_gAMA(ppng, pinf, gamma); }
    if (debug) { jspng_dump_info(stderr, __FUNCTION__, "after set_gAMA", ppng, pinf); }
    
    /* Write the PNG header: */
    png_write_info (ppng, pinf);
    if (debug) { jspng_dump_info(stderr, __FUNCTION__, "after write_info", ppng, pinf); }
    fflush(wr);

    /* -- OUTPUT TRANSFORMATION REQUESTS -------------------------------- */

    /* Ask the PNG library to squeeze input samples into 1,2,4-bit images if they fit: */
    if (bits < 8) { png_set_packing(ppng); } 
    if (debug) { jspng_dump_info(stderr, __FUNCTION__, "after set_packing", ppng, pinf); }

    /* ------------------------------------------------------------------ */

    /* Samples per row in PNM and PNG images: */
    int spr = (int)(cols*chns); /* Samples per row. */

    /* Assumes non-interlaced PNG: */
    /* Allocate one output scanline buffer: */
    unsigned int png_row_bytes = (unsigned int)png_get_rowbytes(ppng, pinf); /* Bytes per row in file. */
    unsigned int usr_samples_bytes = (bits <= 8 ? 1 : 2); /* Bytes per pixel for {png_write_row} */
    unsigned int usr_row_bytes = spr * usr_samples_bytes; /* Bytes per row for {png_write_row} */
    png_byte *png_rowP = notnull(malloc(usr_row_bytes + 1), "no mem");
    if (debug) { fprintf(stderr, "png_row_bytes = %u usr_row_bytes = %u", png_row_bytes, usr_row_bytes); }
        
    /* Process rows: */
    int row;
    for (row = 0; row < rows; ++row)
      {
        if (debug) { fprintf(stderr, "    [ "); }
        png_bytep png_P = png_rowP;
        pnm_sample_t *pnm_P = img->smp[row];
        int k;
        for (k = 0; k < spr; k++)
          { int smp = (*pnm_P); pnm_P++;
            /* Apply scale factor: */
            if (imaxval != omaxval) { smp = (int)floor(scale*smp+0.5); }
            /* Store in PNG row buffer: */
            if (debug) { fprintf(stderr, " %d", smp); }
            if (bits <= 8)
              { /* {png_write_row} wants one byte per sample: */
                (*png_P) = (png_byte)smp; png_P++;
              }
            else
              { /* {png_write_row} wants two bytes per sample: */
                (*png_P) = (png_byte)(smp >> 8); png_P++;
                (*png_P) = (png_byte)(smp & 255); png_P++;
              }
          }
        png_write_row(ppng, png_rowP);
        if (debug) { fprintf(stderr, " ]\n"); }
      }    
    png_write_end(ppng, pinf);
    /* png_destroy_write_struct(&ppng, &pinf); */
    png_free_data(ppng, pinf, PNG_FREE_ALL, -1);
    free(png_rowP);
  }

void jspng_dump_info(FILE *wr, const char *func, char *label, png_structp ppng, png_infop pinf)
  { 
    fprintf(wr, "------------------------------------------------------------\n");
    fprintf(wr, "PNG file info - %s - %s\n", func, label);
    fprintf(wr, "image_width = %u\n", (unsigned int)png_get_image_width(ppng, pinf));
    fprintf(wr, "image_height = %u\n", (unsigned int)png_get_image_height(ppng, pinf));
    fprintf(wr, "channels = %u\n", (unsigned int)png_get_channels(ppng, pinf));
    fprintf(wr, "bit_depth = %d\n", png_get_bit_depth (ppng, pinf));
    /* fprintf(wr, "usr_bit_depth = %d\n", (int)ppng->usr_bit_depth; */
    fprintf(wr, "color_type = %d\n", png_get_color_type(ppng, pinf));
    fprintf(wr, "interlace_type = %d\n", png_get_interlace_type(ppng, pinf));
    fprintf(wr, "compression_type = %d\n", png_get_compression_type(ppng, pinf));
    fprintf(wr, "filter_type = %d\n", png_get_filter_type(ppng, pinf));
    
    png_color_8p sBIT; 
    if (png_get_sBIT(ppng, pinf, &sBIT)) 
      { fprintf(wr, "sBIT = %16p = ( %u %u %u %u %u )\n", (void *)sBIT, sBIT->red, sBIT->green, sBIT->blue, sBIT->gray, sBIT->alpha); }
    else
      { fprintf(wr, "sBIT not specified\n"); }
    
    double gamma;
    if (png_get_gAMA(ppng, pinf, &gamma))
      { fprintf(wr, "gAMA = %25.16e\n", gamma); }
    else
      { fprintf(wr, "gAMA not specified\n"); }
    
    fprintf(wr, "channels = %d\n", png_get_channels(ppng, pinf));
    fprintf(wr, "rowbytes = %u\n", (unsigned int)png_get_rowbytes(ppng, pinf));
    png_bytep sg = (png_bytep)png_get_signature(ppng, pinf);
    if (sg != NULL)
      { fprintf(wr, "signature = %16p = %c%c%c%c%c%c%c%c\n", (void *)sg, sg[0], sg[1], sg[2], sg[3], sg[4], sg[5], sg[6], sg[7]); }
    else
      { fprintf(wr, "png_get_signature failed\n"); }
    /* fprintf(wr, " = %d\n", png_get_(ppng, pinf)); */
    fprintf(wr, "------------------------------------------------------------\n");

  }