#include <stdio.h>
#include <zlib.h>
#include <assert.h>
#include <errno.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <sodium.h>
#include "crc_util.h"
#include "CRCLib.h"
#include "crc.h"

#define CHUNK 1024
const long idat_signature = 1229209940;
const long iend_signature = 1229278788;

FILE *zlib_decompress_data(unsigned char *data_chunk, size_t file_length, char *buff) {
	int ret;
	unsigned int have;
	z_stream strm;
	unsigned char out[CHUNK];
	unsigned char in[CHUNK];
	size_t sz;

	errno=0;
	FILE *data_stream = fmemopen(data_chunk, file_length, "r");
	FILE *of = NULL;
	of = open_memstream(&buff, &sz);
	if(data_stream == NULL) {
		perror("F MEM OPEN");
	}

	strm.zalloc = Z_NULL;
	strm.zfree = Z_NULL;
	strm.opaque = Z_NULL;
	strm.avail_in = 0;
	strm.next_in = Z_NULL;
	ret = inflateInit(&strm);
	if(ret != Z_OK) {
		return NULL;
	}

	do {
		strm.avail_in = fread(in, 1, CHUNK, data_stream);
		if(ferror(data_stream)) {
			(void)inflateEnd(&strm);
			return NULL;
		}
		if(strm.avail_in == 0) {
			break;
		}
		strm.next_in = in;
	
		do {
			strm.avail_out = CHUNK;
			strm.next_out = out;
			ret = inflate(&strm, Z_NO_FLUSH);
			assert(ret != Z_STREAM_ERROR);
			switch(ret) {
				case Z_NEED_DICT:
					ret = Z_DATA_ERROR;
				case Z_DATA_ERROR:
				case Z_MEM_ERROR:
					(void)inflateEnd(&strm);
					printf("Error: %d\n", ret);
					printf("MSG: %s\n", (char*)strm.msg);
					return NULL;
			}
			have = CHUNK - strm.avail_out;
			fwrite(out, 1, have, of);
		} while(strm.avail_out == 0);
	} while(ret != Z_STREAM_END);

	(void)inflateEnd(&strm);

	fclose(data_stream);

	return of;
}

FILE *zlib_compress_data(unsigned char *data_chunk, size_t file_length, char *buff) {
	int ret, flush;
	unsigned int have;
	z_stream strm;
	unsigned char in[CHUNK];
	unsigned char out[CHUNK];
	int level = 9;
	int method = Z_DEFLATED;
	int windowBits = 10;
	int memLevel = 9;
	//int strategy = Z_DEFAULT_STRATEGY;
	int strategy = Z_FILTERED;
	size_t sz;

	FILE *data_stream = fmemopen(data_chunk, file_length, "r");
	FILE *out_data_stream = NULL;
	out_data_stream = open_memstream(&buff, &sz);

	strm.zalloc = Z_NULL;
	strm.zfree = Z_NULL;
	strm.opaque = Z_NULL;
	ret = deflateInit2(&strm, level, method, windowBits, memLevel, strategy);
	if (ret != Z_OK) {
		return NULL;
	}

	do {
		strm.avail_in = fread(in, 1, CHUNK, data_stream);
		if (ferror(data_stream)) {
			(void)deflateEnd(&strm);
			return NULL;
		}
		flush = feof(data_stream) ? Z_FINISH : Z_NO_FLUSH;
		strm.next_in = in;

		do {
			strm.avail_out = CHUNK;
			strm.next_out = out;

			ret = deflate(&strm, flush);
			assert(ret != Z_STREAM_ERROR);
			have = CHUNK - strm.avail_out;
			if(fwrite(out, 1, have, out_data_stream) != have || ferror(out_data_stream)) {
				(void)deflateEnd(&strm);
				return NULL;
			}
		} while(strm.avail_out == 0);
		assert(strm.avail_in == 0);

	} while(flush != Z_FINISH);
	assert(ret == Z_STREAM_END);

	fclose(data_stream);

	(void)deflateEnd(&strm);
	return out_data_stream;
}

unsigned long first_idat(unsigned char *addr) {
	int idat_found = 0;
	unsigned long offset = 8;
	int jump_offset = 0;
	int 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 iend_found = 0;
	int found_idat = 0;
	unsigned long offset = 8;
	int jump_offset = 0;
	int header_type = 0;
	while(iend_found == 0) {
		jump_offset = check_header_length(addr, offset);
		header_type = check_header_length(addr, offset+4);
		if(header_type == iend_signature) {
			iend_found = 1;
		} else {
			if(header_type == idat_signature) {
				found_idat++;
			}
			offset = offset + jump_offset + 12;
		}
	}
	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;
	}
}

void random_data_change(unsigned char *color_data, int width, int length) {
	int searching = 1;
	size_t rounds = 0;
	width = 16;
	int color_range = 3;	
	unsigned char temp_color_data[length];

	memcpy(temp_color_data, color_data, length);

	do {
		rounds++;
		// Creating temporary data set
		memcpy(temp_color_data, color_data, length);
		// Generating random byte to change
		int random_num = randombytes_uniform(length);
		// Checking for index break
		if(random_num % ((width * color_range) + 1)) {
			if(color_data[random_num] == 255) {
				temp_color_data[random_num]--;
			} else {
				temp_color_data[random_num]++;
			}
			char *check_data_buff = NULL;
			FILE *check_data = zlib_compress_data(temp_color_data, length, check_data_buff);
			size_t check_data_length = 0;

			unsigned char *check_data_array = file_to_char_array(check_data, &check_data_length);

			fclose(check_data);
			free(check_data_buff);

			unsigned char full_data[check_data_length+4];
			full_data[0] = 0x49;
			full_data[1] = 0x44;
			full_data[2] = 0x41;
			full_data[3] = 0x54;
			for(int i = 0; i < check_data_length; i++) {
				full_data[i+4] = check_data_array[i];
			}
			unsigned int temp_crc = crc(full_data, check_data_length);
			if ((temp_crc >> (8*3)) == 10 ) {
				printf("Found in %zu rounds!\n", rounds);
				searching = 0;
			}
			free(check_data_array);

		}
	
	} while(searching == 1);
}

int change_idat_content(unsigned char *addr, char *message, int accuracy, unsigned long offset) {
	//printf("Starting IDAT Tranform\n");
	if(accuracy > 4) {
		printf("Warning, accuracy cannot be larger than 4");
		return EXIT_FAILURE;
	}
	if(accuracy > 2) {
		printf("Notice, this could take a long time...");
	}
	if(total_idat(addr) < strlen(message)) {
		printf("Warning, message exceeds IDAT amount");
		return EXIT_FAILURE;
	}

	int idat_length = check_header_length(addr, offset);
	printf("IDAT Length: %d\n", idat_length);

	int prop_found = 0;
	long size = 1;
	long rounds = 0;
	size_t idat_byte_length = 0;
	unsigned char* idat_data = calloc(size, sizeof(unsigned char));
	for(size_t i = 0; i <= idat_length; i++) {
		if(i == size) {
			size *= 2;
			idat_data = reallocarray(idat_data, size, sizeof(unsigned char));
		}
		idat_data[i] = addr[i+offset+8];
		idat_byte_length = i;
	}
	unsigned char temp_idat_data[idat_byte_length];
	while(prop_found == 0) {
		for(int i = 0; i <= idat_length; i++) {
			temp_idat_data[i] = idat_data[i];
		}
		//int r = randombytes_uniform(5) + 1;
		//int j = randombytes_uniform(idat_length);
		
		//temp_idat_data[j] = (temp_idat_data[j] + r) % 255;
		//unsigned char crc_check[size+4];
		//unsigned int idat_header[] = { 0x49, 0x44, 0x41, 0x54 };
		//for(int i = 0; i < 4; i++) {
		//	crc_check[i] = idat_header[i];
		//}
		//for(int i = 0; i < idat_byte_length; i++) {
		//	crc_check[i] = temp_idat_data[i+4];
		//}
		//unsigned int crcnum = crc(crc_check, idat_length);
		//unsigned int checked_crc = crcnum >> (8*3);
		//rounds++;
		//if(checked_crc == 61) {
			// Setting TEMP IDAT DATA BACK TO ORIGINAL
			// TO STOP DECOMPRESSION CORRUPTION
			//temp_idat_data[j] = addr[offset+8+j];
			// Decompressing Data
			char *uncom_data_buff = NULL;
			FILE *uncom_data = zlib_decompress_data(temp_idat_data, idat_byte_length, uncom_data_buff);
			size_t uncom_data_size = 0;

			unsigned char *uncom_data_array = file_to_char_array(uncom_data, &uncom_data_size);

			random_data_change(uncom_data_array, 16, uncom_data_size);

			fclose(uncom_data);
			free(uncom_data_buff);

			//printf("Found %d in %d rounds\n", checked_crc ,rounds);
			//printf("Full CRC: %08X\n", crcnum);
			//printf("Original: %02X\n", idat_data[j]);
			//printf("Change offset: %d to hex: %02X\n", j, temp_idat_data[j]);
			//addr[offset+8+j] = temp_idat_data[j];
			//update_file_crc(addr, offset, crcnum);
			prop_found = 1;
		//}
	}

	return 0;
}

int main() {
	FILE *fp;
	size_t i = 0;
	unsigned long offset = 0;
	unsigned char *message = malloc(sizeof(char));
	message[0] = '\0';

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

	fp = fopen("./1.png", "rt");
	unsigned char *file_data = file_to_char_array(fp, &i);
	fclose(fp);

	offset = first_idat(file_data);
	change_idat_content(file_data, message, 1, offset);
	//create_cc_file(file_data, i);
}