#include <X11/Xlib.h>
#include <Imlib2.h>
#include <stdio.h>
#include <math.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <fcntl.h>

// Image data, read/write data, byte size of data, r/g/b channel state
void write_steg(DATA32 **out, unsigned char *data, int size, int *state);
void read_steg(DATA32 **out, unsigned char *data, int size, int *state);

// color * == DATA32 *
typedef struct {
   unsigned char blue;
   unsigned char green;
   unsigned char red;
   unsigned char a;
} color;

int main(int argc, char **argv)
{
   int width, height;
   Imlib_Image image;
   int fd, i, n, state;
   struct stat buf;
   unsigned char buffer[4096];
   DATA32 *in, *out;

   if (argc < 4 || argc > 5) {
      fprintf(stderr, "Usage: %s (i <image> <data> <outimage> | o <image> <data>)\n", argv[0]);
      return 1;
   } else if (argv[1][0] == 'i') {
      if (stat(argv[2], &buf) < 0) {
         fprintf(stderr, "Error with %s.\n", argv[2]);
	 exit(1);
      }
      if ((fd = open(argv[3], O_RDONLY)) < 0) {
         fprintf(stderr, "Error with %s.\n", argv[2]);
	 exit(1);
      }
      fstat(fd, &buf);
      image = imlib_load_image(argv[2]);
      imlib_context_set_image(image);
      in = imlib_image_get_data_for_reading_only();
      width = imlib_image_get_width();
      height = imlib_image_get_height();
      if (width*height*3/8 < buf.st_size+4) {
         fprintf(stderr, "%s isn't big enough for %s.  %s has %d space but needs %d.\n", argv[2], argv[3], argv[2], width*height*3/8, buf.st_size+4);
	 exit(1);
      }
      imlib_context_set_image(imlib_create_image(width, height));
      out = imlib_image_get_data();

      // Copy image, so the result can be written in a different format.
      for (i = 0; i < width*height; i++)
         out[i] = in[i];

      // Start writing on red.
      state = 0;
      write_steg(&out, (unsigned char *) (&buf.st_size), 4, &state);

      n = buf.st_size;
      while (n && ((i = read(fd, &buffer, 4096)) >= 0))
      {
        write_steg(&out, buffer, i, &state);
	n -= i;
      }
      imlib_save_image(argv[4]); 
   } else if (argv[1][0] == 'o') {
      if (stat(argv[2], &buf) < 0) {
         fprintf(stderr, "Error with %s.\n", argv[2]);
	 exit(1);
      }
      image = imlib_load_image(argv[2]);
      imlib_context_set_image(image);
      in = imlib_image_get_data_for_reading_only();
      width = imlib_image_get_width();
      height = imlib_image_get_height();
      if (width*height*3/8 < 4) {
         fprintf(stderr, "This image is way too small to hold anything.\n");
	 exit(1);
      }
      // Start reading on red.
      state = 0;
      read_steg(&in, (unsigned char *) &n, 4, &state);

      if (width*height*3/8-4 < n) {
         fprintf(stderr, "The supplied size in the image s too big.  Giving up.\n");
	 exit(1);
      }
      
      fd = creat(argv[3], 0600);
      while (n)
      {
        i = (n > 4096) ? 4096 : n;
        read_steg(&in, buffer, i, &state);
        write(fd, buffer, i);
	n -= i;
      }
   } else {
      fprintf(stderr, "Usage: %s (i <image> <data> <outimage> | o <image> <data>)\n", argv[0]);
      return 1;
   }
   return 0;
}

void write_steg(DATA32 **out, unsigned char *data, int size, int *state)
{
   int i, j;
  
   if (size > 0) {
      for (j = 0; j < size; j++)
        for (i = 0; i < 8; i++)
	{
	  switch (*state)
	  {
	     case 0:
	     ((color *) *out)->red = (((color *) *out)->red & 254)
	                           | ((data[j]>>i) & 1);
	     (*state)++;
	     break;
	     case 1:
	     ((color *) *out)->green = (((color *) *out)->green & 254)
	                             | ((data[j]>>i) & 1);
	     (*state)++;
	     break;
	     case 2:
	     ((color *) *out)->blue = (((color *) *out)->blue & 254)
	                            | ((data[j]>>i) & 1);
	     *state = 0;
	     (*out)++;
	     break;
	  }
	}
   }
}

void read_steg(DATA32 **out, unsigned char *data, int size, int *state)
{
   int i, j;
  
   if (size > 0) {
      for (j = 0; j < size; j++)
        for (i = 0; i < 8; i++)
	{
	  switch (*state)
	  {
	     case 0:
	     data[j] = data[j] & ~(1<<i);
	     data[j] = data[j] | ((((color *) *out)->red & 1) << i);
	     (*state)++;
	     break;
	     case 1:
	     data[j] = data[j] & ~(1<<i);
	     data[j] = data[j] | ((((color *) *out)->green & 1) << i);
	     (*state)++;
	     break;
	     case 2:
	     data[j] = data[j] & ~(1<<i);
	     data[j] = data[j] | ((((color *) *out)->blue & 1) << i);
	     *state = 0;
	     (*out)++;
	     break;
	  }
	}
   }
}