#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <math.h>
#include <inttypes.h>

#include <zlib.h>
#include <zconf.h>

#define CISO_MAGIC 0x4F534943 // "CISO"
#define ZISO_MAGIC 0x4F53495A // "ZISO"
#define MB64_IN_BYTES 64 * 1024 * 1024

unsigned char header_pvd[6] = {0x01, 0x43, 0x44, 0x30, 0x30, 0x31}; // 0x01 + "CD001"
unsigned char magic_psp[8] = {0x50, 0x53, 0x50, 0x20, 0x47, 0x41, 0x4D, 0x45}; // "PSP GAME"
unsigned char record_icon0[10] = {0x09, 0x49, 0x43, 0x4F, 0x4E, 0x30, 0x2E, 0x50, 0x4E, 0x47}; // 0x09 + "ICON0.PNG"

int *file_iso;
FILE *ptr_input;
int magic;
int alignment_index;
int size_block;
uint32_t *buffer_index;
uint8_t *buffer_output;
uint8_t *buffer_input;

void error_and_exit(char *error) {
	printf(error);
	exit(-1);
}

void decompress(uint64_t total_block) {
	int size_cmp = 0;
	int counter = 0;
	
	for(uint64_t block = 0; block < total_block; block++) {
		uint32_t index = buffer_index[block];
		int plain = index & 0x80000000;
		index = index & 0x7fffffff;
		uint64_t pos_read = index << (alignment_index);
		
		uint64_t size_read;
		if(plain) size_read = size_block;
		else {
			uint32_t index_2 = buffer_index[block + 1] & 0x7fffffff;
			if(alignment_index) size_read = (index_2 - index + 1) << alignment_index;
			else size_read = (index_2 - index) << (alignment_index);
		}
		fseek(ptr_input, pos_read, SEEK_SET);
		z_stream infstream;
		infstream.zalloc = Z_NULL;
		infstream.zfree = Z_NULL;
		infstream.opaque = Z_NULL;
		infstream.avail_in = fread(buffer_input, 1, size_read, ptr_input);
		
		if(plain) {
			memcpy(buffer_output, buffer_input, size_read);
			size_cmp = size_read;
		} else {
			if(magic == CISO_MAGIC) {
				infstream.next_in = buffer_input;
				infstream.avail_out = size_block;
				infstream.next_out = buffer_output;

				inflateInit2(&infstream, -MAX_WBITS);
				inflate(&infstream, Z_NO_FLUSH);
				inflateEnd(&infstream);
				
				size_cmp = size_block - infstream.avail_out;
			} else if(magic == ZISO_MAGIC) error_and_exit("ZISO format not supported yet, sorry\n");
		}
		memcpy(file_iso + (counter / sizeof(uint32_t)), buffer_output, size_cmp);
		counter += size_cmp;
	}
	rewind(ptr_input);
}

int main(int argc, char **argv) {
	if(argc != 3) error_and_exit("USAGE: <inputfile> <outputfile>\n");

	ptr_input = fopen(argv[1], "r");
	if (ptr_input == NULL) error_and_exit("Error opening input file\n");
	FILE *ptr_output = fopen(argv[2], "wb");
	if (ptr_output == NULL) error_and_exit("Error opening output file\n");
	
	fseek(ptr_input, 0L, SEEK_END);
	uint64_t size_file = ftell(ptr_input);
	rewind(ptr_input);
	if(size_file < 32) error_and_exit("Input file too small to be valid\n");

	// TODO: will this break on different endians?
	unsigned char header_cso[0x18];
	fread(header_cso, 1, sizeof(header_cso), ptr_input);
	
	magic = header_cso[0x0] + (header_cso[0x1] << 8) + (header_cso[0x2] << 16) + (header_cso[0x3] << 24); 
	if(magic != CISO_MAGIC && magic != ZISO_MAGIC) error_and_exit("File couldn't be identified as CISO or ZISO\n");
	
	uint64_t size_uncompressed = 0; 
	for(int i = 7; i >= 0; --i) { 
		size_uncompressed <<= 8; 
		size_uncompressed |= (uint64_t)header_cso[0x8 + i]; 
	}
	size_block = header_cso[0x10] + (header_cso[0x11] << 8) + (header_cso[0x12] << 16) + (header_cso[0x13] << 24);
	int version = header_cso[0x14];
	alignment_index = header_cso[0x15];
	if((magic == ZISO_MAGIC && version != 1) || !size_uncompressed || !size_block) error_and_exit("Corrupt header in specified ZISO file\n");
	
	int total_block = floor(size_uncompressed / size_block);
	int limit64mb_block = floor(MB64_IN_BYTES / size_block);
	if(total_block > limit64mb_block) total_block = limit64mb_block; // First 64MB should be enough to get the location of the file, we can decrypt more if needed later
	int size_index = (total_block + 1) * sizeof(uint32_t);
	
	buffer_index = calloc(1, size_index);
	buffer_output = calloc(1, size_block);
	buffer_input = calloc(1, size_block * 2);
	file_iso = (int*) calloc(1, size_uncompressed * sizeof(int));
	if(!buffer_index || !buffer_output || !buffer_input || !file_iso) error_and_exit("Couldn't allocate enough memory\n");
	fread(buffer_index, 1, size_index, ptr_input);
	decompress(total_block);
	
	uint64_t *location_pvd = memmem(file_iso, size_uncompressed * sizeof(int), header_pvd, sizeof(header_pvd));
	unsigned char pvd[2048] = {0};
	unsigned char magic_iso[8] = {0};
	unsigned char first18bytes[18] = {0};
	memcpy(&pvd, location_pvd, 2048);
	for(int i = 0; i < 8; i++) magic_iso[i] = pvd[8 + i];
	for(int i = 0; i < 8; i++) if(magic_iso[i] != magic_psp[i]) error_and_exit("File couldn't be identified as PSP ISO\n");
	
	uint64_t *location_record_icon0_name = memmem(file_iso, size_uncompressed * sizeof(int), record_icon0, sizeof(record_icon0));
	uint64_t *location_record_icon0 = location_record_icon0_name - 4;
	memcpy(&first18bytes, location_record_icon0, 18);
	int location = first18bytes[2] + (first18bytes[3] << 8) + (first18bytes[4] << 16) + (first18bytes[5] << 24);
	location = location * size_block;
	int size = first18bytes[10] + (first18bytes[11] << 8) + (first18bytes[12] << 16) + (first18bytes[13] << 24);
	int *png = (int*) malloc(size * sizeof(int));
	
	if((location + size + (2048 * 8)) > MB64_IN_BYTES) { // I forgot why I add (2048 * 8) since my first prototype two years ago, sorry...
		free(buffer_index);
		free(file_iso);
		
		total_block = floor((location + size + (2048 * 8)) / size_block);
		size_index = (total_block + 1) * sizeof(uint32_t);
		
		memset(buffer_output, 0, size_block);
		memset(buffer_input, 0, size_block * 2);
		buffer_index = calloc(1, size_index);
		file_iso = (int*) calloc(1, (location + size + (2048 * 8)) * sizeof(int));

		if(!buffer_index || !buffer_output || !buffer_input || !file_iso) error_and_exit("Couldn't allocate enough memory\n");
		fread(header_cso, 1, sizeof(header_cso), ptr_input);
		fread(buffer_index, 1, size_index, ptr_input);
		decompress(total_block);
	}

	memcpy(png, (file_iso + (location / sizeof(int))), size);
	fwrite(png, 1, size, ptr_output);
	fclose(ptr_output);
	fclose(ptr_input);	
	free(buffer_index);
	free(buffer_output);
	free(buffer_input);
}