/* +-------------------------------------------------------------------+ */
/* | Copyright 1996-2000, Greg Roelofs (newt@pobox.com)                | */
/* | Last revised:  26 August 2000                                     | */
/* +-------------------------------------------------------------------+ */

/* $Id: readWritePNG.c,v 1.1 2000/09/02 22:09:06 torsten Exp $ */

#include <stdio.h>
#include <stdlib.h>
#include "png.h"
#include "image.h"
#include "rwTable.h"

#ifndef TRUE
#  define TRUE 1
#  define FALSE 0
#endif

#ifndef Trace
#  ifdef XPAINT_PNG_DEBUG
#    define Trace(x)   fprintf x ; fflush(stderr)
#  else
#    define Trace(x)
#  endif
#endif

typedef struct _jmpbuf_wrapper {
  jmp_buf jmpbuf;
} jmpbuf_wrapper;

extern void compressColormap(Image *image);
static void xpaint_png_error_handler(png_structp png_ptr, png_const_charp msg);

/* this can be shared between reading/writing code since they never overlap: */
static jmpbuf_wrapper xpaint_jmpbuf_struct;


int
TestPNG(char *file)  /* gets called a LOT on the first image:  brushes? */
{
    char header[8];
    FILE *fp = fopen(file, "rb");   /* libpng requires ANSI; so do we */

    if (!fp)
        return 0;

    fread(header, 1, 8, fp);
    fclose(fp);

    return png_check_sig(header, 8);
}



Image *
ReadPNG(char *file)
{
    FILE          *fp;
    png_structp   png_ptr;
    png_infop     info_ptr;
    png_uint_32   width, height;
    png_colorp    palette;
    int           bit_depth, color_type, interlace_type, num_palette, rowbytes;
    int           i, hasAlpha=FALSE;
    int           level, npasses;
    Image         *image = NULL;


    Trace((stderr, "\nGRR ReadPNG:  reading file %s\n", file));
    if ((fp = fopen(file, "rb")) == (FILE *)NULL) {
        RWSetMsg("Error opening input file");
        return NULL;
    }

    png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING,
      &xpaint_jmpbuf_struct, xpaint_png_error_handler, NULL);
    if (!png_ptr) {
        RWSetMsg("Error allocating PNG png_ptr memory");
        fclose(fp);
        return NULL;
    }

    info_ptr = png_create_info_struct(png_ptr);
    if (!info_ptr) {
        RWSetMsg("Error allocating PNG info_ptr memory");
        png_destroy_read_struct(&png_ptr, NULL, NULL);
        fclose(fp);
        return NULL;
    }

    if (setjmp(xpaint_jmpbuf_struct.jmpbuf)) {
        RWSetMsg("Fatal libpng error (longjmp() called while reading)");
        fprintf(stderr,
          "Fatal libpng error (longjmp() called while reading)\n");
        fflush(stderr);
        png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
        fclose(fp);
        return NULL;
    }

    png_init_io(png_ptr, fp);
    png_read_info(png_ptr, info_ptr);  /* read all PNG info up to image data */

    png_get_IHDR(png_ptr, info_ptr, &width, &height, &bit_depth, &color_type,
      &interlace_type, NULL, NULL);
    Trace((stderr, "GRR ReadPNG:  width = %d, height = %d\n", width, height));

    switch (color_type) {

        case PNG_COLOR_TYPE_PALETTE:
            Trace((stderr, "GRR ReadPNG:  PNG_COLOR_TYPE_PALETTE\n"));
            if (!png_get_PLTE(png_ptr, info_ptr, &palette, &num_palette)) {
                RWSetMsg("Error:  PLTE chunk not found in palette image");
                fprintf(stderr,
                  "ReadPNG error:  PLTE chunk not found in palette image\n");
                fflush(stderr);
                png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
                fclose(fp);
                return NULL;
            }
            image = ImageNewCmap(width, height, num_palette);
            for (i = 0;  i < num_palette;  ++i)
                ImageSetCmap(image, i, palette[i].red, palette[i].green,
                  palette[i].blue);
            /* GRR:  still need to get image data into `image' */
            break;

        case PNG_COLOR_TYPE_RGB:
            Trace((stderr, "GRR ReadPNG:  PNG_COLOR_TYPE_RGB\n"));
            if (bit_depth == 16) {
                RWSetMsg("Stripping 48-bit RGB image to 24 bits");
                fprintf(stderr,
                  "ReadPNG:  stripping 48-bit RGB image to 24 bits\n");
                fflush(stderr);
                png_set_strip_16(png_ptr);
            }
            image = ImageNew(width, height);
            break;

        case PNG_COLOR_TYPE_GRAY:   /* treat grayscale as special colormap */
            Trace((stderr, "GRR ReadPNG:  PNG_COLOR_TYPE_GRAY\n"));
            if (bit_depth == 16) {
                RWSetMsg("Stripping 16-bit grayscale image to 8 bits");
                fprintf(stderr,
                  "ReadPNG:  stripping 16-bit grayscale image to 8 bits\n");
                fflush(stderr);
                png_set_strip_16(png_ptr);
                bit_depth = 8;
            }
            image = ImageNewCmap(width, height, 1 << bit_depth);
            switch (bit_depth) {
                case 1:
                    image->isBW = TRUE;
                    ImageSetCmap(image, 0,   0,   0,   0);
                    ImageSetCmap(image, 1, 255, 255, 255);
                    break;
                case 2:
                    image->isGrey = TRUE;
                    ImageSetCmap(image, 0,   0,   0,   0);
                    ImageSetCmap(image, 1,  85,  85,  85);  /* 255/3 */
                    ImageSetCmap(image, 2, 170, 170, 170);
                    ImageSetCmap(image, 3, 255, 255, 255);
                    break;
                case 4:
                    image->isGrey = TRUE;
                    for (i = 0;  i < 16;  ++i) {
                        level = i * 17;  /* 255/15 */
                        ImageSetCmap(image, i, level, level, level);
                    }
                    break;
                case 8:
                    image->isGrey = TRUE;
                    for (i = 0;  i < 256;  ++i)
                        ImageSetCmap(image, i, i, i, i);
                    break;
            }
            break;

        case PNG_COLOR_TYPE_RGB_ALPHA:
            Trace((stderr, "GRR ReadPNG:  PNG_COLOR_TYPE_RGB_ALPHA\n"));
            if (bit_depth == 16) {
                RWSetMsg("Stripping 64-bit RGBA image to 32 bits");
                fprintf(stderr,
                  "ReadPNG:  stripping 64-bit RGBA image to 32 bits\n");
                fflush(stderr);
                png_set_strip_16(png_ptr);
            }
            /* need to split alpha and RGB/gray the hard way, sigh */
            hasAlpha = TRUE;
            image = ImageNew(width, height);
            ImageMakeMask(image);   /* maskData is always 1 byte deep */
            break;

        case PNG_COLOR_TYPE_GRAY_ALPHA:
            Trace((stderr, "GRR ReadPNG:  PNG_COLOR_TYPE_GRAY_ALPHA\n"));
            if (bit_depth == 16) {
                RWSetMsg("Stripping 32-bit gray+alpha image to 16 bits");
                fprintf(stderr,
                  "ReadPNG:  stripping 32-bit gray+alpha image to 16 bits\n");
                fflush(stderr);
                png_set_strip_16(png_ptr);
            }
            hasAlpha = TRUE;
            image = ImageNewGrey(width, height);
            ImageMakeMask(image);   /* maskData is always 1 byte deep */
            break;

        default:
            fprintf(stderr, "ReadPNG:  unknown image type (%d)\n", color_type);
            fflush(stderr);
            RWSetMsg("Unknown PNG image type");
            fprintf(stderr, "ReadPNG error:  unknown PNG image type\n");
            fflush(stderr);
            png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
            fclose(fp);
            return NULL;
    }

    if (bit_depth < 8)
        png_set_packing(png_ptr);

    if (interlace_type)
        npasses = png_set_interlace_handling(png_ptr);

    png_read_update_info(png_ptr, info_ptr);

    rowbytes = png_get_rowbytes(png_ptr, info_ptr);

    if (hasAlpha) {
        png_bytep *row_pointers, png_data;

	if ((row_pointers =
             (png_bytep *)malloc(height * sizeof(png_bytep))) == 0)
        {
	    RWSetMsg("ReadPNG:  can't allocate memory for row_pointers (1)");
            fprintf(stderr,
              "ReadPNG:  can't allocate memory for row_pointers (1)\n");
            fflush(stderr);
            ImageDelete(image);
            png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
            fclose(fp);
	    return NULL;
	}
	
        /* very inefficient (memory-wise) to allocate entire image again, but
         * no easy way around it:  libpng returns the image and alpha channel
         * interspersed, and interlaced alpha images just make matters worse
         */
        png_data = (png_bytep)malloc(height*rowbytes);
        if (!png_data) {
            RWSetMsg("Unable to allocate temporary storage for alpha image");
            fprintf(stderr, "ReadPNG:  unable to malloc png_data\n");
            fflush(stderr);
            ImageDelete(image);
            png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
            fclose(fp);
            return NULL;
        }

        /* only bit depths of 8 and 16 support alpha channels */
        for (i = 0;  i < height;  ++i)
            row_pointers[i] = (png_bytep)png_data + i*rowbytes;

        png_read_image(png_ptr, row_pointers);

        if (color_type == PNG_COLOR_TYPE_RGB_ALPHA) {
            png_bytep png=png_data, rgb=image->data, alpha=image->maskData;

            for (i = 0;  i < height*width;  ++i) {
                *rgb++   = *png++;
                *rgb++   = *png++;
                *rgb++   = *png++;
                *alpha++ = *png++;  /* GRR:  invert?  doesn't seem to be used */
            }
        } else {  /* PNG_COLOR_TYPE_GRAY_ALPHA */
            png_bytep png=png_data, gray=image->data, alpha=image->maskData;

            for (i = 0;  i < height*width;  ++i) {
                *gray++  = *png++;
                *alpha++ = *png++;  /* GRR:  invert? */
            }
        }
        free(png_data);   /* whew */

    } else {  /* no alpha channel */
        png_bytep *row_pointers;

	if ((row_pointers =
             (png_bytep *)malloc(height * sizeof(png_bytep))) == 0)
        {
	    RWSetMsg("ReadPNG:  can't allocate memory for row_pointers (2)");
            fprintf(stderr,
              "ReadPNG:  can't allocate memory for row_pointers (2)\n");
            fflush(stderr);
            ImageDelete(image);
            png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
            fclose(fp);
	    return NULL;
	}
        for (i = 0;  i < height;  ++i)
            row_pointers[i] = (png_bytep)image->data + i*rowbytes;

        png_read_image(png_ptr, row_pointers);

    } /* end if (hasAlpha) */

    /* GRR:  ideally should read all other info (text chunks, etc.) and save
     *       whatever ones make sense for the possible output file...but we
     *       don't, both because we're too darned lazy to bother and because
     *       XPaint is too simple-minded to be able to use it
     */

    png_read_end(png_ptr, info_ptr);
    png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
    fclose(fp);

    return image;

} /* end function ReadPNG() */



static int
WritePNG(char *file, Image *image, int interlace_type)
{
    int           i;
    int           bit_depth, color_type, num_palette, rowbytes;
    png_colorp    palette;
    png_structp   png_ptr;
    png_infop     info_ptr;
    png_textp     software;
    FILE          *fp;


    Trace((stderr, "\nGRR WritePNG:  %s\n", file));
    Trace((stderr, "GRR WritePNG:  %d x %d, scale = %d\n",
      image->width, image->height, image->scale));

    if (!(fp = fopen(file, "wb")))
        return 1;

    png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING,
      &xpaint_jmpbuf_struct, xpaint_png_error_handler, NULL);
    if (!png_ptr) {
        fclose(fp);
        return 1;
    }

    info_ptr = png_create_info_struct(png_ptr);
    if (!info_ptr) {
        png_destroy_write_struct(&png_ptr, NULL);
        fclose(fp);
        return 1;
    }

    if (setjmp(xpaint_jmpbuf_struct.jmpbuf)) {
        RWSetMsg("Fatal libpng error (longjmp() called while writing)");
        fprintf(stderr,
          "Fatal libpng error (longjmp() called while writing)\n");
        fflush(stderr);
        png_destroy_write_struct(&png_ptr, &info_ptr);
        fclose(fp);
        return 1;
    }

    png_init_io(png_ptr, fp);

    if (image->isBW) {
        if (image->maskData) {
            color_type = PNG_COLOR_TYPE_GRAY_ALPHA;
            bit_depth = 8;   /* promote to full grayscale */
        } else {
            color_type = PNG_COLOR_TYPE_GRAY;
            bit_depth = 1;
        }
        Trace((stderr, "GRR WritePNG:  B/W, bit_depth = %d\n", bit_depth));

    } else if (image->isGrey) {
        color_type = image->maskData? PNG_COLOR_TYPE_GRAY_ALPHA :
                                      PNG_COLOR_TYPE_GRAY;
        if (image->cmapPacked)
            bit_depth = 8;
        else {
            Trace((stderr,
              "GRR WritePNG:  isGrey, cmapSize = %d (before compressing), ",
              image->cmapSize));
            compressColormap(image);
            Trace((stderr, "%d (after)\n", image->cmapSize));
            if (image->cmapSize > 16)
                bit_depth = 8;
            else if (image->cmapSize > 4)
                bit_depth = 4;
            else if (image->cmapSize > 2)
                bit_depth = 2;
            else
                bit_depth = 1;
            Trace((stderr, "GRR WritePNG:  isGrey, picked bit_depth = %d\n",
              bit_depth));
        }

    } else if (image->scale == 3) {
        Image *cmapImage;

        /* try compressing image to palette mode, but don't force if too big */
        if (!image->maskData)   /* can't store alpha mask with palette image */
            cmapImage = ImageCompress(image, 256, 1);

        if (cmapImage) {
            image = cmapImage;  /* original was deleted in ImageCompress() */
        } else {
            color_type = image->maskData? PNG_COLOR_TYPE_RGB_ALPHA :
                                          PNG_COLOR_TYPE_RGB;
            bit_depth = 8;
            Trace((stderr, "GRR WritePNG:  RGB, bit_depth = 8\n"));
        }
    }

    /* either we're on an 8-bit or smaller display, or image->scale was 3 and
     * ImageCompress() worked
     */
    if (image->scale == 1) {
        color_type = PNG_COLOR_TYPE_PALETTE;
        if (image->maskData) {
            fprintf(stderr,
              "WritePNG:  can't use alpha mask with colormapped image\n");
            fflush(stderr);
        }
        Trace((stderr, "GRR WritePNG:  palette, bit_depth = "));
        if (!image->cmapPacked)
            compressColormap(image);
        if (image->cmapSize > 16)
            bit_depth = 8;
        else if (image->cmapSize > 4)
            bit_depth = 4;
        else if (image->cmapSize > 2)
            bit_depth = 2;
        else
            bit_depth = 1;
        palette = (png_colorp)image->cmapData;   /* GRR:  seems to work... */
#if 0
0000000000000000000
        for (i = 0;  i < image->cmapSize;  ++i) {
            palette[i].red = image->cmapData ...
            palette[i].green = 
            palette[i].blue = 
        }
#endif
        num_palette = image->cmapSize;
        png_set_PLTE(png_ptr, info_ptr, palette, num_palette);
        Trace((stderr, "%d, num_palette = %d\n", bit_depth, num_palette));
    }

    Trace((stderr, "GRR WritePNG:  image->scale = %d, image->width = %d, image->height = %d\n", image->scale, image->width, image->height));

    png_set_IHDR(png_ptr, info_ptr, image->width, image->height, bit_depth,
      color_type, interlace_type, PNG_COMPRESSION_TYPE_DEFAULT,
      PNG_FILTER_TYPE_DEFAULT);

#if 0	/* if we don't know what the gamma is, don't guess */
    if (xpaint_foo->gamma > 0.0)
        png_set_gAMA(png_ptr, info_ptr, xpaint_foo->gamma);
#endif

    /* store the modification time, at least */
    {
        png_time  modtime;

        png_convert_from_time_t(&modtime, time(NULL));
        png_set_tIME(png_ptr, info_ptr, &modtime);
    }

    /* only one text comment:  Software */
    {
        png_text textdata;
        char software_text[40];

        sprintf(software_text, "XPaint %s", XPAINT_VERSION);
        textdata.compression = PNG_TEXT_COMPRESSION_NONE;
        textdata.key = "Software";
        textdata.text = software_text;
        png_set_text(png_ptr, info_ptr, &textdata, 1);
    }

    /* write all chunks up to (but not including) first IDAT */
    Trace((stderr, "GRR WritePNG:  calling png_write_info()\n"));
    png_write_info(png_ptr, info_ptr);
    png_write_flush(png_ptr);

    /* any transformations must be set *after* png_write_info() is called */
    Trace((stderr, "GRR WritePNG:  calling png_set_packing()\n"));
    png_set_packing(png_ptr);   /* squish multiple pixels into each byte */
/*  png_set_compression_level(png_ptr, Z_BEST_COMPRESSION);  */

    if (image->maskData) {
        /* alpha channel version */
        fprintf(stderr, "WritePNG:  sorry, can't write alpha images yet\n");
        fflush(stderr);
        png_destroy_write_struct(&png_ptr, &info_ptr);
        fclose(fp);
        return 1;
    } else {
        int rowbytes = image->scale * image->width;
        png_bytep *row_pointers;

	if ((row_pointers =
	     (png_bytep *)malloc(image->height * sizeof(png_bytep))) == 0)
        {
	    RWSetMsg("WritePNG:  can't allocate memory for row_pointers");
            fprintf(stderr,
              "WritePNG:  can't allocate memory for row_pointers\n");
            fflush(stderr);
            png_destroy_write_struct(&png_ptr, &info_ptr);
            fclose(fp);
	    return 1;
	}

        for (i = 0;  i < image->height;  ++i)
            row_pointers[i] = (png_bytep)image->data + i*rowbytes;

#ifdef XPAINT_PNG_DEBUG
        {
            int wrapcount = (image->scale == 1)? 20 : 24;
            int rownum = (image->height >> 1);
            png_bytep row = row_pointers[rownum];

            Trace((stderr, "GRR WritePNG:  data for row %d:", rownum));
            for (i = 0;  i < rowbytes;  ++i, ++row) {
                if (i % wrapcount == 0)
                    Trace((stderr, "\n  "));
                Trace((stderr, "%02d ", *row));
            }
            Trace((stderr, "\nGRR WritePNG:  done with row %d.\n", rownum));
        }
#endif

        png_write_image(png_ptr, row_pointers);
    }

    png_write_end(png_ptr, NULL);

    png_destroy_write_struct(&png_ptr, &info_ptr);
    fclose(fp);

    return 0;

} /* end function WritePNG() */



int
WritePNGn(char *file, Image *image)
{
    return WritePNG(file, image, 0);
}



int
WritePNGi(char *file, Image *image)
{
    return WritePNG(file, image, 1);
}



static void
xpaint_png_error_handler (png_structp png_ptr, png_const_charp msg)
{
  jmpbuf_wrapper  *jmpbuf_ptr;

  /* this function, aside from the extra step of retrieving the "error
   * pointer" (below) and the fact that it exists within the application
   * rather than within libpng, is essentially identical to libpng's
   * default error handler.  The second point is critical:  since both
   * setjmp() and longjmp() are called from the same code, they are
   * guaranteed to have compatible notions of how big a jmp_buf is,
   * regardless of whether _BSD_SOURCE or anything else has (or has not)
   * been defined. */

  fprintf(stderr, "XPaint:  fatal libpng error: %s\n", msg);
  fflush(stderr);

  jmpbuf_ptr = png_get_error_ptr(png_ptr);
  if (jmpbuf_ptr == NULL) {         /* we are completely hosed now */
    fprintf(stderr, "XPaint:  "
      "EXTREMELY fatal libpng error: jmpbuf unrecoverable; terminating.\n");
    fflush(stderr);
    exit(99);
  }

  longjmp(jmpbuf_ptr->jmpbuf, 1);
}