#pragma GCC optimize("O0")

#include <endian.h>
#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <sodium.h>
#include <getopt.h>
#include <pthread.h>
#include <math.h>
#include "crc_util.h"
#include "CRCLib.h"
#include "pspng.h"
#include "compress_util.h"

//#define MAX_IDAT_SIZE 16384
#define MAX_WINDOW_SIZE 32768
#define MAX_IDAT_SIZE 8192

const long idat_signature = 1229209940;
const long iend_signature = 1229278788;

unsigned long first_idat(unsigned char *addr) {
	int idat_found = 0;
	unsigned long offset = 8;
	long jump_offset = 0;
	unsigned long header_type = 0;
	while(idat_found == 0) {
		jump_offset = check_header_length(addr, offset);
		header_type = check_header_length(addr, offset+4);
		if(header_type == idat_signature) {
			idat_found = 1;
		} else {
			offset = offset + jump_offset + 12;
		}
	}
	return offset;
}

int total_idat(unsigned char *addr) {
	int searching = 0;
	int found_idat = 0;
	unsigned long offset = 0;
	unsigned long idat_length = 0;
	unsigned long header_type = 0;

	unsigned long first_idat_offset = first_idat(addr);
	offset = first_idat_offset;

	while(searching == 0) {
		found_idat++;
		idat_length = check_header_length(addr, offset);
		// Jumping Offset + IDAT LENGTH + 3 bytes for 2 CRCs and 1 LEN fields
		offset = offset+12+idat_length;
		header_type = check_header_length(addr, offset+4);
		if(header_type != idat_signature) {
			searching = 1;
		}
	}
	return found_idat;
}

int update_file_crc(unsigned char *addr, unsigned long offset , unsigned int crc_num) {
	int startCRC = 8 + offset + check_header_length(addr, offset);
	unsigned char new_crc;
	for(int i = 0; i < 4; i++) {
		new_crc = crc_num >> (8*(3-i)) & 0xFF;
		addr[startCRC+i] = new_crc;
	}
	return 0;
}

static unsigned char* create_raw_pixel_values(unsigned char *data, int width, int color_range, size_t length) {
	// Generating twice the size we migth need
	length = length * 4;
	//length = 500000;
	unsigned char* raw_data = calloc(length, sizeof(unsigned char*));
	size_t filter_chunk_byte = 0;

	for(size_t i = 0; i < length; i++) {
		filter_chunk_byte = i  - (i % ((width * color_range) + 1));
		// Grabbing raw pixel data
		if(data[filter_chunk_byte] == 0) {
			// Filter: None
			raw_data[i] = data[i];
		} else if(data[filter_chunk_byte] == 1) {
			// Filter: Sub
			if ((i % ((width * color_range) + 1)) == 0) {
				// Setting filter type
				raw_data[i] = data[i];
			} else if ((i - filter_chunk_byte) <= color_range) {
				// Setting initial color_range byte
				raw_data[i] = data[i];
			} else {
				// Setting remaining bytes
				raw_data[i] = data[i] + raw_data[i - color_range];
			}
		
		} else if(data[filter_chunk_byte] == 3) {
			// Filter: Average
			if((i % ((width * color_range) + 1)) == 0) {
				// Setting filter type
				raw_data[i] = data[i];
			} else {
				// Setting remaining bytes
				if(filter_chunk_byte == 0) {
					// If no scanline is present before current
					if((i - filter_chunk_byte) <= color_range) {
						// If no pixel present before current
						raw_data[i] = data[i] + floor((0 + 0) / 2);
					} else {
						// If pixel present before current
						raw_data[i] = data[i] + floor((data[i - color_range] + 0) / 2);
					}
				
				} else {
					// Scanline is present before current
					if((i - filter_chunk_byte) <= color_range) {
						// If no pixel present before current
						raw_data[i] = data[i] + floor((0 + raw_data[i - ((width * color_range) + 1)]) / 2);
					} else {
						// If pixel present before current
						raw_data[i] = data[i] + floor((raw_data[i - color_range] + raw_data[i - ((width * color_range) + 1)]) / 2);
					}
				}
			}
		
		} else {
			// Filter: Error Unsupported
			printf("Exiting due to unsupported filter type\n");
			exit(EXIT_FAILURE);
		}
	}

	return raw_data;
}

static void random_window_bit_change(unsigned char *data, unsigned char *raw_data, int width, int rounds, int color_range, size_t length, size_t offset) {
	size_t random_num = 0;
	size_t filter_chunk_byte = 0;
	for (int i = 0; i < rounds; i++) {
		random_num = 0;
		do {
			random_num = randombytes_uniform(length) + offset;
			filter_chunk_byte = (random_num - (random_num % ((width * color_range) + 1)));
		} while((random_num % ((width * color_range) + 1)) == 0);

		if(data[filter_chunk_byte] == 0) {
			// None Filter Type Change
			if(data[random_num] == 255) {
				data[random_num]--;
				raw_data[random_num]--;
			} else {
				data[random_num]++;
				raw_data[random_num]++;
			}
		} else if(data[filter_chunk_byte] == 1) {
			// Sub Filter Type Change
			if(data[random_num] == 255) {
				data[random_num]--;
				raw_data[random_num]--;
			} else {
				data[random_num]++;
				raw_data[random_num]++;
			}
			// Fixing Next Pixel in Scanline
			if(!(random_num+3 > (filter_chunk_byte + (width * color_range) + 1))){
				data[random_num+color_range] = raw_data[random_num+color_range] - raw_data[random_num];
			}
			// Checking Next Scanline type
			if(data[filter_chunk_byte+((width * color_range) + 1)] == 3){
				data[random_num+((width * color_range)+1)] = raw_data[random_num+((width * color_range)+1)] - floor((raw_data[random_num+((width * color_range)+1)-color_range] + raw_data[random_num]) / 2);
			}
		} else if(data[filter_chunk_byte] == 2) {
			// Add Support
			if(data[random_num] == 255) {
				data[random_num]--;
			} else {
				data[random_num]++;
			}
		} else if(data[filter_chunk_byte] == 3) {
			// Average Filter Type Change
			if(raw_data[random_num] == 255) {
				raw_data[random_num]--;
			} else {
				raw_data[random_num]++;
			}
			// Update Current Pixel
			data[random_num] = raw_data[random_num] - floor((raw_data[random_num-color_range] + raw_data[random_num-((width * color_range) + 1)]) / 2);
			// Fixing Next Pixel in Scanline
			if(!(random_num+color_range > (filter_chunk_byte + (width * color_range) + 1))){
				data[random_num+color_range] = raw_data[random_num+color_range] - floor((raw_data[random_num] + raw_data[random_num-((width * color_range)+ 1)+color_range]) / 2);
			}
			// Checking Next Scanline type
			if(data[filter_chunk_byte+((width * color_range) + 1)] == 3){
				data[random_num+((width * color_range)+1)] = raw_data[random_num+((width * color_range)+1)] - floor((raw_data[random_num+((width * color_range)+1)-color_range] + raw_data[random_num]) / 2);
			}
		} else if(data[filter_chunk_byte] == 4) {
			// Add Support
			if(data[random_num] == 255) {
				data[random_num]--;
			} else {
				data[random_num]++;
			}
		} else {
			printf("ERROR ;( %ld\n", filter_chunk_byte);
		}

	}
	return;
}

static int verify_crc_chunks(unsigned char *data, size_t data_length, int crc_depth, unsigned int max_data_length, char message, size_t iteration) {
	unsigned char* testing_chunk;
	unsigned int crc_check = 0;
	unsigned int test_data_length = 0;
	unsigned int crc_check_length;
	int rounds = 0;

	for(int i = 0; i < crc_depth; i++) {
		if((max_data_length * (i+1)) > data_length) {
			test_data_length = (data_length - (max_data_length * i));
			testing_chunk = calloc(test_data_length+4, sizeof(unsigned char));
			// If this is met we need to stop verification
			i = crc_depth;
		} else {
			test_data_length = max_data_length;
			testing_chunk = calloc(test_data_length+4, sizeof(unsigned char));
		}
		// Setting first chunk to IDAT
		testing_chunk[0] = 0x49;
		testing_chunk[1] = 0x44;
		testing_chunk[2] = 0x41;
		testing_chunk[3] = 0x54;
		for(int j = 0; j < test_data_length; j++) {
			testing_chunk[j+4] = data[(test_data_length*rounds)+j+(MAX_WINDOW_SIZE*iteration)];
		}
		crc_check_length = test_data_length+4;
		crc_check = crc(testing_chunk, crc_check_length);

		if ((crc_check >> (8*3)) != message ) {
			free(testing_chunk);
			return 1;
		}
		printf("FOUND: %c\n", message);
		// Used as an alternative to i
		// Since i needs to be changed it max data length is met
		rounds++;
	}
	free(testing_chunk);
	return 0;
}

static int crc_embed_data(unsigned char *data, unsigned int data_length, int bit_width, int color_range, unsigned int sliding_window, char message, size_t offset , size_t iteration, size_t prior_offset) {
	unsigned char *check_data;
	unsigned char *raw_data;
	unsigned char *new_data = calloc(1, sizeof(unsigned char));
	size_t check_data_length = 0;
	int match_crc = 1;
	
	raw_data = create_raw_pixel_values(data, bit_width, color_range, sliding_window);
	random_window_bit_change(data, raw_data, bit_width, 2, color_range, sliding_window, offset);

	if(offset > 500000) {
		new_data = calloc(300000, sizeof(unsigned char));
		for(int k = 0; k < 300000; k++) {
			new_data[k] = data[offset-(offset-prior_offset)+k];
		}
		zlib_compress_data(new_data, 110000, &check_data, &check_data_length);
		match_crc = verify_crc_chunks(check_data, check_data_length, 1, MAX_IDAT_SIZE, message, 1);
	} else {
		zlib_compress_data(data, data_length, &check_data, &check_data_length);
		match_crc = verify_crc_chunks(check_data, check_data_length, 1, MAX_IDAT_SIZE, message, iteration);
	}

	if(match_crc == 0) {
		printf("COM SIZE: %zu\n", check_data_length);
		free(new_data);
		free(check_data);
		free(raw_data);
		return 0;
	}
	free(new_data);
	free(check_data);
	free(raw_data);
	return 1;
}

void *random_data_change_thread_call(void *w) {
	struct EMBED_THREAD_STRUCT *data = w;
	unsigned char* temp_color_data = calloc(data->uncom_data_len, sizeof(unsigned char));
	char cur_message;
	int searching = 1;
	do {
		// Creating temporary data set
		memcpy(temp_color_data, data->data, data->uncom_data_len);

		// Look into testing these values more for speed
		if(data->uncom_data_len > 800000) {
			data->data_len = 90000 + (90000 * data->cur_iteration);
		} else {
			data->data_len = data->uncom_data_len;
		}
		if(data->uncom_data_len > 36000) {
			data->win_size = 36000 + (36000 * data->cur_iteration);
		} else {
			data->win_size = data->uncom_data_len;
		}

		cur_message = data->message[data->cur_iteration];

		searching = crc_embed_data(temp_color_data, data->data_len, data->width, data->color_range, data->win_size, cur_message, data->cur_offset, data->cur_iteration, data->offset[data->cur_iteration-1]);

		pthread_mutex_lock(&data->mutex_lock);
		if (searching == 0 && *data->searching == 1) {
			*data->searching = searching;
			printf("CUR MESSAGE: %c\n", cur_message);
			memcpy(data->data, temp_color_data, data->uncom_data_len);
			pthread_mutex_unlock(&data->mutex_lock);
			break;
		}
		pthread_mutex_unlock(&data->mutex_lock);

	} while(*data->searching == 1);
	free(temp_color_data);
	pthread_exit(0);
}

size_t generate_offset(unsigned char *data, size_t data_len, size_t iteration) {
	unsigned char *com_data_buff = NULL;
	unsigned char *uncom_data_buff = NULL;
	size_t com_data_size = 0;
	size_t uncom_data_size = 0;

	printf("Gen Offset\n");
	// See if we can make this call faster
	zlib_compress_data(data, data_len, &com_data_buff, &com_data_size);

	zlib_decompress_data(com_data_buff, (MAX_WINDOW_SIZE * (iteration+1)), &uncom_data_buff, &uncom_data_size);

	printf("NEW SIZE: %zu\n", uncom_data_size);

	free(com_data_buff);
	free(uncom_data_buff);
	return uncom_data_size;
}

void random_data_change(unsigned char *color_data, unsigned char *width, size_t length, char *message) {
	int searching = 1;
	int core_count = 28;
	unsigned int sliding_window = 0;
	unsigned int compress_data_length = 0;
	// Needs to be turned into a variable
	int color_range = 3;
	int waiting = 0;
	// Union for width type cast
	union {
		uint32_t width_int;
		unsigned char width_array[4];
	}w;
	memcpy(w.width_array, width, 4);

	struct EMBED_THREAD_STRUCT *t_data = malloc(sizeof(struct EMBED_THREAD_STRUCT));

	if(pthread_mutex_init(&t_data->mutex_lock, NULL) != 0) {
		printf("Mutex Lock Error\n");
		return;
	}
	// Building Struct Data
	t_data->searching = &searching;
	t_data->message = message;
	t_data->data = color_data;
	t_data->data_len = compress_data_length;
	t_data->uncom_data_len = length;
	// Original Offset set to 0
	t_data->cur_offset = 0;
	t_data->width = be32toh(w.width_int);
	t_data->color_range = color_range;
	t_data->win_size = sliding_window;
	pthread_t tid;
	t_data->thread_id = tid;
	for(int j = 0; j < strlen(message); j++) {
		// Setting Iteration bit and searching
		t_data->cur_iteration = j;
		*t_data->searching = 1;
		for(int i = 0; i < core_count; i++) {
			pthread_create(&tid, NULL, random_data_change_thread_call, t_data);
		}
		// Waiting for all threads to complete
		waiting = 0;
		do {
			pthread_join(tid, NULL);
			waiting++;
		} while(waiting!=core_count);
		pthread_mutex_destroy(&t_data->mutex_lock);

		// Only generate new offset if not last char in message
		if(j != (strlen(message) - 1)) {
			// Changing offset check size to not be full length of picture to speed up generation
			//t_data->cur_offset = generate_offset(color_data, t_data->uncom_data_len, j);
			t_data->cur_offset = generate_offset(color_data, (200000 * (j+1)), j);
			t_data->offset[j] = t_data->cur_offset;
		}
	}
	return;
}

void build_png_file(struct PNG_FILE_STRUCT *png_file, char *out_file_name) {
	FILE *fp;
	unsigned char *full_data;
	unsigned int int_crc;
	unsigned char new_crc[4];
	unsigned char IEND_DATA[12] = { 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82};

	fp = fopen(out_file_name, "w");

	union{
		unsigned char data[sizeof(struct PNG_START_FILE_STRUCT)];
		struct PNG_START_FILE_STRUCT png_data;
	}start_data;

	start_data.png_data = png_file->png_start_data;

	printf("IHDR SIZE: %zu\n", sizeof(struct PNG_START_FILE_STRUCT));

	// IHDR Data
	for(int i = 0; i < sizeof(start_data.data); i++) {
		fputc(start_data.data[i], fp);
	}

	// Generating IDAT CHUNKS
	int idat_loop = 0;
	uint32_t current_len = 0;
	do {
		// Setting IDAT length chunk variable
		if ((be32toh(png_file->png_idat_data.idat_data_length)-(MAX_IDAT_SIZE*idat_loop)) > MAX_IDAT_SIZE) {
			current_len = MAX_IDAT_SIZE;
		} else {
			current_len = (be32toh(png_file->png_idat_data.idat_data_length)-(MAX_IDAT_SIZE*idat_loop));
		}

		// IDAT LENGTH WRITE
		for(int i = 0; i < 4; i++) {
			//fputc(png_file->png_idat_data.idat_length[i], fp);
			fputc(current_len >> (8*(3-i)), fp);
		}
		// IDAT HEADER WRITE
		for(int i = 0; i < 4; i++) {
			fputc(png_file->png_idat_data.idat_header[i], fp);
		}
		// IDAT DATA WRITE
		//for(int i = 0; i < be32toh(png_file->png_idat_data.idat_data_length); i++) {
		for(int i = 0; i < current_len; i++) {
			fputc(png_file->png_idat_data.idat_data[i+(MAX_IDAT_SIZE*idat_loop)], fp);
		}

		// Generating CRC
		//unsigned char *full_data = malloc(be32toh(png_file->png_idat_data.idat_data_length)+4);
		full_data = malloc(current_len+4);
		for(int i = 0; i < 4; i++) {
			full_data[i] = png_file->png_idat_data.idat_header[i];
		}
		for(int i = 0; i < current_len; i++) {
			full_data[i+4] = png_file->png_idat_data.idat_data[i+(MAX_IDAT_SIZE*idat_loop)];
		}
	
		int_crc = crc(full_data, current_len+4);
	
		// IDAT CRC WRITE
		for(int i = 0; i < 4; i++) {
			new_crc[i] = int_crc >> (8*(3-i)) & 0xFF;
			fputc(new_crc[i], fp);
		}

		// Adding to loop count
		idat_loop++;
	} while((idat_loop-1) < (be32toh(png_file->png_idat_data.idat_data_length) / MAX_IDAT_SIZE));

	// IEND Data
	for(int i = 0; i < 12; i++) {
		fputc(IEND_DATA[i], fp);
	}
	free(full_data);
	fclose(fp);
}

unsigned char* populate_idat_array(unsigned char *addr, unsigned long offset, size_t *idat_byte_length) {
	size_t size = 8;
	//size_t idat_byte_length = 0;
	size_t idat_length = 0;
	//unsigned char* idat_data = calloc(size, sizeof(unsigned char));
	unsigned char* idat_data = calloc(size, sizeof(unsigned char));

	for(size_t i = 0; i < total_idat(addr); i++) {
		idat_length = check_header_length(addr, offset);
		for(size_t j = 0; j < idat_length; j++) {
			if(*idat_byte_length == size) {
				size *= 2;
				idat_data = reallocarray(idat_data, size, sizeof(unsigned char));
			}
			idat_data[*idat_byte_length] = addr[offset+8+j];
			*idat_byte_length += 1;
		}
		// Offset 3 bytes for 2 length fields and 1 crc
		offset = offset + idat_length + 12;
	}
	return idat_data;
}

int change_idat_content(unsigned char *addr, struct PNG_FILE_STRUCT *png_file, char *message, int accuracy, unsigned long offset, char *out_file_name) {
	if(accuracy > 4) {
		// Accuracy is currently unused
		printf("Warning, accuracy cannot be larger than 4");
		return EXIT_FAILURE;
	} else if (accuracy > 2) {
		printf("Notice, this could take a long time...");
	}
	// Checking if total IDAT / 4 (Since Sliding Window is assumed to be 32k)
	if((total_idat(addr) / 4) < strlen(message)) {
		printf("Warning, message exceeds IDAT amount\n");
		exit(EXIT_FAILURE);
	}

	size_t idat_byte_length = 0;
	unsigned char *idat_data = populate_idat_array(addr, offset, &idat_byte_length);

	// Decompressing Data
	unsigned char *uncom_data_buff = NULL;
	size_t uncom_data_size = 0;
	zlib_decompress_data(idat_data, idat_byte_length, &uncom_data_buff, &uncom_data_size);

	printf("ORIG UNCOM LEN: %zu\n", uncom_data_size);

	// Start data testing
	random_data_change(uncom_data_buff, png_file->png_start_data.file_width, uncom_data_size, message);
	
	// Compress Data
	unsigned char *com_data_buff;
	size_t com_data_size = 0;
	zlib_compress_data(uncom_data_buff, uncom_data_size, &com_data_buff, &com_data_size);

	printf("COM LEN: %zu\n", com_data_size);

	png_file->png_idat_data.idat_data = calloc(com_data_size, sizeof(unsigned char));

	png_file->png_idat_data.idat_data_length = be32toh(com_data_size);

	for(size_t i = 0; i < com_data_size; i++) {
		png_file->png_idat_data.idat_data[i] = com_data_buff[i];
	}

	// Build PNG File
	build_png_file(png_file, out_file_name);

	// Freeing used memory
	free(uncom_data_buff);
	free(com_data_buff);
	free(idat_data);
	free(png_file->png_idat_data.idat_data);

	return 0;
}

// This is where it all starts
int main(int argc, char **argv) {
	FILE *fp;
	size_t i = 0;
	unsigned long offset = 0;
	int uncompress_call = 0;
	int compress_call = 0;
	struct PNG_FILE_STRUCT png_file_data;
	char *in_file_name = NULL;
	char *out_file_name = NULL;
	char *message = NULL;

	static const struct option long_options[] = {
		{"help", no_argument, NULL, 'h'},
		{"file", required_argument, NULL, 'f'},
		{"outfile", required_argument, NULL, 'o'},
		{"message", required_argument, NULL, 'm'},
		{"compress", no_argument, NULL, 'c'},
		{"uncompress", no_argument, NULL, 'u'},
		{0, 0, 0, 0}
	};

	const char* usage =
		"Usage: pspng [options]\n\n"
		"  -c  --compress       	Compress Message\n"
		"  -u  --uncompress     	Currently used for debugging\n"
		"  -h, --help			Shows help message\n"
		"  -f, --file			Denotes input file\n"
		"  -o, --outfile		Denotes output file\n"
		"  -m, --message		Encoded message\n"
		"\n";

	int c;
	while (1) {
		int option_index = 0;
		c = getopt_long(argc, argv, "hf:o:m:uc", long_options ,&option_index);
		if(c == -1) {
			break;
		}
		switch(c) {
			case 'h':
				printf("%s", usage);
				exit(EXIT_SUCCESS);
			case 'f':
				in_file_name = optarg;
				break;
			case 'o':
				out_file_name = optarg;
				break;
			case 'm':
				message = optarg;
				break;
			case 'c':
				compress_call = 1;
				break;
			case 'u':
				uncompress_call = 1;
				break;
		}
	}

	if(in_file_name == NULL) {
		printf("Input file required!\n");
		exit(EXIT_FAILURE);
	}
	
	if(compress_call == 1) { 
		if(out_file_name == NULL) {
			printf("Output file required!\n");
			exit(EXIT_FAILURE);
		} else if(message == NULL) {
			printf("Message required!\n");
			exit(EXIT_FAILURE);
		}
	}

	if(sodium_init() == -1) {
		return EXIT_FAILURE;
	}

	fp = fopen(in_file_name, "rt");
	if (fp == NULL) {
		printf("File error\n");
		exit(EXIT_FAILURE);
	}


	unsigned char *file_data = file_to_char_array(fp, &i);
	fclose(fp);

	populate_start_png(file_data, &png_file_data.png_start_data);

	offset = first_idat(file_data);

	if(uncompress_call == 1) {
		size_t idat_byte_length = 0;
		unsigned long cur_idat_len = 0;
		unsigned char* idat_data = populate_idat_array(file_data, offset, &idat_byte_length);

		printf("Decompressed Data:\n");

		do {
			for(int j = 0; j < 1020; j++) {
				cur_idat_len = 0;
				for(int i = 0; i < 4; i++) {
					cur_idat_len += (file_data[offset+i] << (24-(8*i)));
				}
					//printf("%c", file_data[offset+cur_idat_len+8]);
				if((j % 4) == 0) {
					if(31 < file_data[offset+cur_idat_len+8] && file_data[offset+cur_idat_len+8] < 128) {
						printf("%c", file_data[offset+cur_idat_len+8]);
					} else {
						offset = idat_byte_length;
						break;
					}
				}
				// Offset Plus idat length + 3 bytes
				offset += cur_idat_len + 12;
			}
			printf("\n");
		} while(offset < idat_byte_length);

		free(idat_data);
		free(file_data);
		return EXIT_SUCCESS;
	}

	if(compress_call == 1) {
		populate_idat_png(file_data, &png_file_data.png_idat_data, offset);
	
		change_idat_content(file_data, &png_file_data, message, 1, offset, out_file_name);
	
		free(file_data);
		return EXIT_SUCCESS;
	}

	free(file_data);
	return EXIT_FAILURE;
}