/* See {jsjpeg_image.h} */
/* Last edited on 2016-04-01 01:27:06 by stolfilocal */

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

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

#include <jpeglib.h>
#undef FALSE
#undef TRUE

#include <affirm.h>
#include <jsfile.h>
#include <jspnm.h>
#include <jspnm_image.h>
#include <jsjpeg_image.h>

pnm_image_t *jsjpeg_image_read (char *name, bool_t verbose, int *kind)
  { FILE *rd = open_read(name, verbose);
    pnm_image_t *img = jsjpeg_image_fread(rd, kind);
    if (rd != stdin) { fclose(rd); }
    return img;
  }

void jsjpeg_image_write (char *name, pnm_image_t *img, int quality, bool_t verbose)
  { FILE *wr = open_write(name, verbose);
    jsjpeg_image_fwrite(wr, img, quality);
    if (wr != stdout) { fclose(wr); }
  }

pnm_image_t *jsjpeg_image_fread (FILE *rd, int *kind)
  {
    /* Create the decompressor structures: */
    struct jpeg_decompress_struct jdec;
    struct jpeg_error_mgr jerm;
    jdec.err = jpeg_std_error(&jerm);
    jpeg_create_decompress(&jdec);
    jpeg_stdio_src(&jdec, rd);
    
    /* Read the file header and start the decompressor: */
    demand(jpeg_read_header(&jdec, TRUE) == 1, "{jpeg_read_header} failed");
    (void) jpeg_start_decompress(&jdec);  /* !!! Should check the return value? !!! */
    
    #define jdemand(cond,msg) \
      if (! (cond)) { jpeg_abort_decompress(&jdec); demand(FALSE, (msg)); }
    
    /* Get the image size: */
    int cols = jdec.output_width;
    int rows = jdec.output_height;
    int chns = jdec.output_components;
    jdemand((chns > 0) && (chns <= 4), "invalid channel count");

    /* Return the image colorspace: */
    (*kind) = jdec.jpeg_color_space;

    /* Alocate the JPEG decompression sample buffer: */    
    int spr = cols * chns;   /* Samples per row in decompressor output. */
    JSAMPARRAY buffer = (*jdec.mem->alloc_sarray)((j_common_ptr)&jdec, JPOOL_IMAGE, spr, 1);

    /* Alocate the in-core image: */    
    pnm_image_t *img = pnm_image_new(cols, rows, chns);
    img->maxval = MAXJSAMPLE;

    /* Read scanline by scanline: */
    while (jdec.output_scanline < rows) {
      int r = jdec.output_scanline;
      unsigned int nread = jpeg_read_scanlines(&jdec, buffer, 1);
      jdemand(nread == 1, "{jpeg_read_scanlines} failed");
      pnm_sample_t *img_row = img->smp[r];
      JSAMPLE *bu_row = buffer[0];
      int k;
      for (k = 0; k < spr; k++) { img_row[k] = (int) bu_row[k]; }
    }
    
    /* Cleanup: */
    (void) jpeg_finish_decompress(&jdec);  /* !!! Should check the return value? !!! */
    jpeg_destroy_decompress(&jdec);

    return img;
    
    #undef jdemand
  }

void jsjpeg_image_fwrite (FILE *wr, pnm_image_t *img, int quality)
  {
    /* !!! Extend to other types besides grayscale and RGB. !!! */

    /* Quality checking: */
    demand((quality > 0) && (quality <= 100), "invalid {quality}");
    
    /* Channel count checking: */
    int chns = img->chns;
    int cols = img->cols;
    int rows = img->rows;
    demand((chns == 1) || (chns == 3), "invalid channel count");
    
    /* Determine the sample scaling factor: */
    if (img->maxval > MAXJSAMPLE)
      { fprintf(stderr, "%s: !! warning:", __FUNCTION__);
        fprintf(stderr, "image {maxval = %d} larger for JPEG encoder's {maxval = %d}\n", img->maxval, MAXJSAMPLE);
        fprintf(stderr, "  samples will be individually downscaled with slight loss of precision.\n");
      }
    else if ((MAXJSAMPLE + 1) % (img->maxval + 1) != 0)
      { fprintf(stderr, "%s: !! warning:", __FUNCTION__);
        fprintf(stderr, "image {maxval = %d} JPEG encoder's {maxval = %d}\n", img->maxval, MAXJSAMPLE);
        fprintf(stderr, "  samples will be individually upscaled with slight loss of precision.\n");
      }
    double scale = ((double)(MAXJSAMPLE + 1))/((double)(img->maxval + 1)); 

    /* Create the compressor structures: */
    struct jpeg_compress_struct jcmp;
    struct jpeg_error_mgr jerm;
    jcmp.err = jpeg_std_error(&jerm);
    jpeg_create_compress(&jcmp);
    jpeg_stdio_dest(&jcmp, wr);

    /* Start the compressor: */
    J_COLOR_SPACE space = (chns == 1 ? JCS_GRAYSCALE : JCS_RGB);
    jcmp.image_width = cols;
    jcmp.image_height = rows;
    jcmp.input_components = chns;
    jcmp.in_color_space = space;
    jpeg_set_defaults(&jcmp);
    /* jpeg_set_colorspace(&jcmp, space); */
    jpeg_set_quality(&jcmp, quality, TRUE);
    jpeg_start_compress(&jcmp, TRUE);
    
    #define jdemand(cond,msg) \
      if (! (cond)) { jpeg_abort_compress(&jcmp); demand(FALSE, (msg)); }

    /* Alocate the JPEG compression sample buffer: */
    int spr = chns * cols; /* Samples per scanline. */
    JSAMPARRAY buffer = (*jcmp.mem->alloc_sarray)((j_common_ptr) &jcmp, JPOOL_IMAGE, spr, 1);

    /* Write scanline by scanline: */
    while (jcmp.next_scanline < rows) {
      int r = jcmp.next_scanline;
      pnm_sample_t *img_row = img->smp[r];
      JSAMPLE *bu_row = buffer[0];
      int k;
      for (k = 0; k < spr; k++) 
        { bu_row[k] = (JSAMPLE)(scale == 1.0 ? img_row[k] : (int)(scale*img_row[k])); }

      unsigned int nwrote = jpeg_write_scanlines(&jcmp, buffer, 1);
      jdemand(nwrote == 1, "{jpeg_read_scanlines} failed");
    }

    /* Cleanup: */
    jpeg_finish_compress(&jcmp);
    jpeg_destroy_compress(&jcmp);
    
    #undef jdemand
  }