#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <stdarg.h>
#include <ctype.h>

#define PNG_DEBUG 3
#include <png.h>

int max = 0;
int x, y;
int width, height;
png_byte color_type;
png_byte bit_depth;
png_structp png_ptr;
png_infop info_ptr;
int number_of_passes;
png_bytep *row_pointers;
png_bytep *row_pointers_hidden;



void error(const char *s, ...)
{
   va_list args;

   va_start(args, s);
   vfprintf(stderr, s, args);
   fprintf(stderr, "\n");
   va_end(args);
   exit(0);
}

void read_png(char *file_name)
{
   unsigned char header[8];

   FILE *fp = fopen(file_name, "rb");

   if(!fp)
      error("[read_png] File %s could not be opened", file_name);

   fread(header, 1, 8, fp);
   if(png_sig_cmp(header, 0, 8))
      error("[read_png] File %s is not recognized as a PNG file", file_name);


   png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);

   if(!png_ptr)
      error("[read_png] png_create_read_struct failed");

   info_ptr = png_create_info_struct(png_ptr);
   if(!info_ptr)
      error("[read_png] png_create_info_struct failed");

   if(setjmp(png_jmpbuf(png_ptr)))
      error("[read_png] Error during init_io");

   png_init_io(png_ptr, fp);
   png_set_sig_bytes(png_ptr, 8);

   png_read_info(png_ptr, info_ptr);

   width = info_ptr->width;
   height = info_ptr->height;
   color_type = info_ptr->color_type;
   bit_depth = info_ptr->bit_depth;

   number_of_passes = png_set_interlace_handling(png_ptr);
   png_read_update_info(png_ptr, info_ptr);


   if(setjmp(png_jmpbuf(png_ptr)))
      error("[read_png] Error during read_image");

   row_pointers = (png_bytep *) malloc(sizeof(png_bytep) * height);

   for(y = 0; y < height; y++)
   {
      row_pointers[y] = (png_byte *) malloc(info_ptr->rowbytes);
   }

   png_read_image(png_ptr, row_pointers);

   fclose(fp);
}


void write_png(char *file_name)
{
   FILE *fp = fopen(file_name, "wb");

   if(!fp)
      error("[write_png] File %s could not be opened for writing",
           file_name);


   png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);

   if(!png_ptr)
      error("[write_png] png_create_write_struct failed");

   info_ptr = png_create_info_struct(png_ptr);
   if(!info_ptr)
      error("[write_png] png_create_info_struct failed");

   if(setjmp(png_jmpbuf(png_ptr)))
      error("[write_png] Error during init_io");

   png_init_io(png_ptr, fp);


   if(setjmp(png_jmpbuf(png_ptr)))
      error("[write_png] Error during writing header");

   png_set_IHDR(png_ptr, info_ptr, width, height,
                bit_depth, color_type, PNG_INTERLACE_NONE,
                PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

   png_write_info(png_ptr, info_ptr);


   if(setjmp(png_jmpbuf(png_ptr)))
      error("[write_png] Error during writing bytes");

   png_write_image(png_ptr, row_pointers);


   if(setjmp(png_jmpbuf(png_ptr)))
      error("[write_png] Error during end of write");

   png_write_end(png_ptr, NULL);

   for(y = 0; y < height; y++)
      free(row_pointers[y]);
   free(row_pointers);

   fclose(fp);
}


int get_value(int x, int y, int rgba)
{
   png_byte *row = row_pointers[y];
   png_byte *ptr = &(row[x * 4]);

   return ptr[rgba];
}

void set_value(int x, int y, int rgba, int value)
{
   png_byte *row = row_pointers[y];
   png_byte *ptr = &(row[x * 4]);

   ptr[rgba] = value;
}


void unhide_file(const char *msg_file)
{
   if(info_ptr->color_type != PNG_COLOR_TYPE_RGBA)
      error("[unhide_file] color_type of input file must be "
           "PNG_COLOR_TYPE_RGBA (is %d)", info_ptr->color_type);

   FILE *f = fopen(msg_file, "w+");
   if(!f)
      error("[unhide_file] fopen()");

   for(y = 0; y < height; y++)
   {
      for(x = 0; x < width; x++)
      {
         if(get_value(x, y, 3)==0)
         {
            char c1 = get_value(x, y, 0);
            char c2 = get_value(x, y, 1);
            char c3 = get_value(x, y, 2);

            if(c1==0 || c2==0 || c3==0)
            {
               fclose(f);
               return;
            }

            fprintf(f, "%c%c%c", c1, c2, c3);
            set_value(x, y, 0, 0);
            set_value(x, y, 1, 0);
            set_value(x, y, 2, 0);
         }
      }
   }

   fclose(f);
}

void hide_file(const char *msg_file)
{
   if(info_ptr->color_type != PNG_COLOR_TYPE_RGBA)
      error("[hide_file] color_type of input file must be "
           "PNG_COLOR_TYPE_RGBA (is %d)", info_ptr->color_type);

   FILE *f = fopen(msg_file, "r");
   if(!f)
      error("[hide_file] fopen()");

   for(y = 0; y < height; y++)
   {
      for(x = 0; x < width; x++)
      {
         if(get_value(x, y, 3)==0)
         {
            char c1=0;
            char c2=0;
            char c3=0;

            if(!feof(f)) c1 = fgetc(f);
            if(!feof(f)) c2 = fgetc(f);
            if(!feof(f)) c3  = fgetc(f);

            set_value(x, y, 0, c1);
            set_value(x, y, 1, c2);
            set_value(x, y, 2, c3);
         }
      }
   }

   fclose(f);
}



int main(int argc, char **argv)
{
   if(argc != 5)
      error("Usage: %s <hide|unhide> <png file in> <png file out> <msg file>",
      argv[0]);
   

   if(strcmp(argv[1], "hide")==0)
   {
      read_png(argv[2]);
      hide_file(argv[4]);
      write_png(argv[3]);
   }
   else if(strcmp(argv[1], "unhide")==0)
   {
      read_png(argv[2]);
      unhide_file(argv[4]);
      write_png(argv[3]);
   }
   else
      error("Usage: %s <hide|unhide> <png file in> <png file out> <msg file>",
      argv[0]);


   return 0;
}