/*
Copyright 2010 Jim Studt. All rights reserved.

Redistribution and use in source and binary forms, with or without modification, are
permitted provided that the following conditions are met:

   1. Redistributions of source code must retain the above copyright notice, this list of
      conditions and the following disclaimer.

   2. Redistributions in binary form must reproduce the above copyright notice, this list
      of conditions and the following disclaimer in the documentation and/or other materials
      provided with the distribution.

THIS SOFTWARE IS PROVIDED BY JIM STUDT ``AS IS'' AND ANY EXPRESS OR IMPLIED
WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> OR
CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

The views and conclusions contained in the software and documentation are those of the
authors and should not be interpreted as representing official policies, either expressed
or implied, of Jim Studt.
*/

#include <stdio.h>
#include <stdarg.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <getopt.h>

static int beVerbose = 0;
static const char *outputName = 0;

struct adds {
    struct adds * next;
    int appSegment;
    int length;     /* needs 2 more for the jpeg encoding */
    unsigned char *content;
} *adds = 0;
    
static const char short_options [] = "h?vi:r:o:";

static const struct option
long_options [] = {
	{ "insert",     required_argument,      NULL,           'i' },
	{ "remove",     required_argument,      NULL,           'r' },
	{ "output",     required_argument,      NULL,           'o' },
        { "help",       no_argument,            NULL,           'h' },
	{ "verbose",    no_argument,            NULL,           'v' },
        { 0, 0, 0, 0 }
};

/*
** Set and index to 1 to skip a given tag.
*/
static unsigned char skip[256] = {0};

//
// A printf to stderr, but only if in verbose mode
//
void verbose( const char *format, ...)
{
    va_list args;

    if ( !beVerbose) return;
    va_start( args, format);
    vfprintf(stderr, format, args);
    va_end(args);
}

//
// A printf to stderr and an exit(1).
//   (adds a trailing newline as well.)
//
void fatal( const char *format, ...) __attribute__ ((noreturn));
void fatal( const char *format, ...)
{
    va_list args;

    va_start( args, format);
    vfprintf(stderr, format, args);
    va_end(args);
    fputc('\n',stderr);

    exit(1);
}


static void usage(FILE *fp, int argc, char **argv)
{
    fprintf (fp,
	     "Usage: %s [options] [input file]\n"
	     " e.g.: %s -i 7=(alpha0:image/png) -i 7=(alpha0:data=)./mask.png image.jpg\n\n"
	     "Options:\n"
	     "  -i | --insert seg:file   insert the contents of file as APPx where x=seg\n"
	     "       --insert seg=string insert the argument as APPx where x=seg\n"
	     "  -r | --remove seg        remove all APPx segments where x=seg\n"
	     "  -o | --output file       send output to file instead of stdout\n"
	     "  -v | --verbose           Print a lot of debug messages\n"
	     "  -h | --help              This help message\n"
	     "",
	     argv[0], argv[0]);
}

static void do_options(int argc, char **argv)
{
    for (;;) {
	int index;
	int c;
                
	c = getopt_long (argc, argv, short_options, long_options, &index);

	if (-1 == c) break;

	switch (c) {
	  case 0: /* getopt_long() flag */
	    break;
	  case 'o':
	    outputName = strdup(optarg);
	    verbose("output will go to %s\n", outputName);
	    break;
	  case 'r':
	      {
		  char *p;
		  int v = strtol( optarg, &p, 10);
		  if ( *p != 0 || v < 0 || v > 15) fatal("Invalid argument to --remove, '%s'", optarg);
		  verbose("will skip APP%d\n", v);
		  skip[(unsigned char)(0xe0 + v)] = 1;
	      }
	      break;
	  case 'i':
	      {
		  int seg;
		  char arg[65536] = "";
		  int contentLength = 0;
		  void *contentPointer = 0;
		  struct adds **tail = &adds;

		  if ( sscanf( optarg, "%d=%65535s", &seg, arg) < 2) {
		      if ( sscanf( optarg, "%d@%65535s", &seg, arg) < 2) {
			  fatal("Poorly formed insertion.");
		      }
		      /* we have a @filename style, get the content */
		      contentPointer = malloc(65536); /* maximum encodable */
		      {
			  int len;
			  FILE *f = fopen( arg, "r");
			  if ( !f) fatal("Unable to open %s for insertion: %s", arg, strerror(errno));
			  len = fread( contentPointer, 1, 65534, f);
			  if ( ferror(f)) fatal("Error reading %s for insertion: %s", arg, strerror(errno));
			  if ( !feof(f)) fatal("Truncated %s for insertion, 65534 bytes maximum.", arg);
			  fclose(f);
			  contentPointer = realloc( contentPointer, len);
			  contentLength = len;
		      }
		  } else {
		      /* we have a literal style */
		      contentPointer = strdup(arg);
		      contentLength = strlen(arg);
		  }
		  if ( seg < 0 || seg > 15) fatal("Segment number of an insert needs to be from 0 to 15.");

		  while( *tail != 0) tail = & (*tail)->next;

		  *tail = (struct adds *)malloc(sizeof(struct adds));
		  (*tail)->next = 0;
		  (*tail)->appSegment = seg;
		  (*tail)->length = contentLength;
		  (*tail)->content = contentPointer;
	      }
	      break;
	  case 'h':
	  case '?':
	    usage (stdout, argc, argv);
	    exit (EXIT_SUCCESS);
	  case 'v':
	    beVerbose++;
	    break;
	  default:
	    usage (stderr, argc, argv);
	    exit (EXIT_FAILURE);
	}
    }
}

static void putOut( int o, FILE *f) {
    if ( fputc( o, f) == EOF) fatal("Failed writing output: %s", strerror(errno));
}

static int doEntropy( FILE *in, FILE *out, int c1, int c2)
{
    int count = 2;

    putOut( c1, out);
    putOut( c2, out);
    /* copy entropy encoded data */
    for (;;) {
	int ch = fgetc(in);
	if ( ch == EOF) fatal("Permature end of JPEG input: %s", strerror(errno));
	
	if ( ch == 0xff) {
	    ch = fgetc(in);
	    if ( ch == EOF) fatal("Permature end of JPEG input: %s", strerror(errno));
	    if ( ch == 0) {
		putOut(0xff,out);
		putOut(0,out);
		count += 2;
	    } else {
		verbose("entropy encoded section, %d bytes\n", count);
		return ch;
	    }
	} else {
	    putOut(ch,out);
	    count++;
	}
    }
}

/*
** We learn the structure of a JPEG from http://www.w3.org/Graphics/JPEG/itu-t81.pdf
** Specifically Appendix B.
** In short we have markers of the from 0xff 0x?? followed by a two byte count, except
** where they are not followed by a two byte count, but instead by an stream of bytes with 
** their 0xff bytes escaped to 0xff 0x00. Oh, and d0-d7, d8, d9, and 01 which don't have any
** content.
*/
void processJPEG( FILE *in, FILE *out)
{
    int shouldEnd = 0;
    int needInsert = (adds != 0);

    for (;;) {
	int tag;
	int ff = fgetc(stdin);
	if ( ff == EOF) {
	    if ( ferror(stdin)) fatal("Failed reading input: %s", strerror(errno));
	    if ( !shouldEnd) fatal("Premature end of JPEG input");
	    break;  /* This is our loop exit <<<<<<<<<<<<< */
	}
	
	tag = fgetc(stdin);
	if ( tag == EOF) fatal("Premature end of JPEG input");
	
	if ( ff != 0xff || tag == 0) tag = doEntropy( in,out, ff, tag);

	if ( (tag >= 0xd0 && tag <= 0xd9) || tag == 0x00) {
	    /* a tag with no content, marker only */
	    verbose("copied xff%02x\n", tag);
	    putOut( 0xff, out);
	    putOut( tag, out);
	    if ( tag == 0xd9) shouldEnd = 1;
	} else {
	    /* copy the marker's block */
	    int msb = fgetc(stdin);
	    int lsb = fgetc(stdin);
	    int len;
	    int skipping = skip[(unsigned char)tag];
	    
	    if ( lsb == EOF || msb == EOF) fatal("Permature end of JPEG input: %s", strerror(errno));
	    len = msb*256+lsb;

	    if ( needInsert && tag == 0xda) {
		/* do our insert right before the SOS tag (start of scan), no particular reason, but legal */
		struct adds *a = adds;

		for ( a = adds; a; a = a->next) {
		    int i;
		    int len = a->length;
		    putOut( 0xff, out);
		    putOut( 0xe0 + a->appSegment, out);
		    len += 2;
		    putOut( (len >> 8) & 0xff, out);
		    putOut( len & 0xff, out);
		    len -= 2;
		    for ( i = 0; i < len; i++) putOut( a->content[i], out);
		    verbose("inserted a APP%d length %d bytes\n", a->appSegment, len);
		}
		needInsert = 0;
	    }

	    if ( !skipping) {
		verbose("copied xff%02x %d bytes\n", tag, len);
		putOut( 0xff, out);
		putOut( tag, out);
		putOut( msb, out);
		putOut( lsb, out);
	    } else {
		verbose("skipped xff%02x %d bytes\n", tag, len);
	    }
	    
	    /* we take 2 off for the len field itself */
	    for ( len -= 2; len > 0; len--) {
		int ch = fgetc(stdin);
		if ( ch == EOF) fatal("Permature end of JPEG input: %s", strerror(errno));
		if ( !skipping) putOut(ch,out);
	    }
	}		
    }
}

int main( int argc, char **argv)
{
    do_options(argc, argv);

    /*
    ** Reopen stdin as required
    */
    switch( argc - optind) {
      case 0:
	/* using stdin */
	break;
      case 1:
	  {
	      const char *inFile = argv[optind];
	      stdin = freopen( inFile, "r", stdin);
	      if ( stdin == 0) fatal("Unable to open %s for reading: %s", inFile, strerror(errno));
	      verbose("input is reopened to %s\n", inFile);
	  }
	break;
      default:
	fatal("Too many arguments, can only process one file.");
    }

    /*
    ** Reopen stdout as required
    */
    if ( outputName) {
	/*
	** I wonder if maybe I should write to a temporary file and rename after success.
	** I don't expect to fail, but that would allow input and output to be the same file.
	** Maybe later.
	*/
	stdout = freopen( outputName, "w", stdout);
	if ( stdout == 0) fatal("Unable to write to %s: %s", outputName, strerror(errno));
	verbose("output is reopened to %s\n", outputName);
    }

    processJPEG( stdin, stdout);

    if ( fclose(stdout) == EOF) fatal("Failed to close output file: %s", strerror(errno));
    if ( fclose(stdin) == EOF) fatal("Failed to close input file: %s", strerror(errno));
    
    return 0;
}