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

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;
	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;
	}
	return 0;
}

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];

	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]++;
			}
			unsigned char *check_data_buff = NULL;
			size_t check_data_length = 0;
			zlib_compress_data(temp_color_data, length, &check_data_buff, &check_data_length);

			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_buff[i];
			}
			unsigned int temp_crc = crc(full_data, check_data_length);
			if ((temp_crc >> (8*3)) == 10 ) {
				printf("Found in %zu rounds!\n", rounds);
				memcpy(color_data, temp_color_data, length);
				searching = 0;
			}
			free(check_data_buff);
		}
	
	} while(searching == 1);
}

void build_png_file(struct PNG_FILE_STRUCT *png_file, char *out_file_name) {
	FILE *fp;

	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;

	// IHDR Data
	for(int i = 0; i < sizeof(start_data.data); i++) {
		fputc(start_data.data[i], fp);
	}
	// IDAT Data
	for(int i = 0; i < 4; i++) {
		fputc(png_file->png_idat_data.idat_length[i], fp);
	}
	for(int i = 0; i < 4; i++) {
		fputc(png_file->png_idat_data.idat_header[i], fp);
	}
	for(int i = 0; i < be32toh(png_file->png_idat_data.idat_data_length); i++) {
		fputc(png_file->png_idat_data.idat_data[i], fp);
	}
	// Generating CRC
	unsigned char full_data[be32toh(png_file->png_idat_data.idat_data_length)+4];
	for(int i = 0; i < 4; i++) {
		full_data[i] = png_file->png_idat_data.idat_header[i];
	}
	for(int i = 0; i < be32toh(png_file->png_idat_data.idat_data_length); i++) {
		full_data[i+4] = png_file->png_idat_data.idat_data[i];
	}

	unsigned int int_crc = crc(full_data, be32toh(png_file->png_idat_data.idat_data_length));
	unsigned char new_crc[4];

	for(int i = 0; i < 4; i++) {
		new_crc[i] = int_crc >> (8*(3-i)) & 0xFF;
		fputc(new_crc[i], fp);
	}
	// IEND Data
	unsigned char IEND_DATA[12] = { 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82};
	for(int i = 0; i < 12; i++) {
		fputc(IEND_DATA[i], fp);
	}
	fclose(fp);
}

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) {
		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((char*)message)) {
		printf("Warning, message exceeds IDAT amount\n");
		exit(EXIT_FAILURE);
	}

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

	long size = 1;
	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];
	for(int i = 0; i <= idat_length; i++) {
		temp_idat_data[i] = idat_data[i];
	}
	// Decompressing Data
	unsigned char *uncom_data_buff = NULL;
	size_t uncom_data_size = 0;
	zlib_decompress_data(temp_idat_data, idat_byte_length, &uncom_data_buff, &uncom_data_size);

	random_data_change(uncom_data_buff, 16, uncom_data_size);
	
	// 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);

	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);

	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;
	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'},
		{0, 0, 0, 0}
	};

	const char* usage =
		"Usage: crc [options]\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:", 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;
		}
	}

	if(in_file_name == NULL) {
		printf("Input file required!\n");
		exit(EXIT_FAILURE);
	} else 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);

	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);
}