#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <dirent.h>
#include <errno.h>
#include <sys/stat.h>

#include "zip/unzip.h"
#include "zip/zip.h"

#include "rar.h"

#ifdef _WIN32
char dir_sep = '\\';
char *path_sep = ";";
char *exe_ext = ".exe";
#else
char dir_sep = '/';
char *path_sep = ":";
char *exe_ext = "";
#endif

static int level = 0;

int unsafe(char character) {
	switch(character) {
		case '"':
		#ifndef _WIN32
		case '\\': case '$':
		#endif
			return 1;
		default:
			return 0;
	}
}

char* escape(char *unescaped) {
	int i, j = 0;
	for(i = 0; i < strlen(unescaped); ++i, ++j) if(unsafe(unescaped[i])) ++j;
	char *ret = malloc(sizeof(char)*(j+1));
	ret[j] = 0;
	j = 0;
	for(i = 0; i < strlen(unescaped); ++i, ++j) {
		if(unsafe(unescaped[i])) {
			ret[j] = '\\';
			++j;
		}
		ret[j] = unescaped[i];
	}
	return ret;
}
int is_jpeg(char *filename) {
	char *filext = strrchr(filename,'.');
	if(!filext) return 0;
	return (!strcasecmp(filext,".jpg") || !strcasecmp(filext,".jpeg") || !strcasecmp(filext,".jpe"));
}

void to_baseline(char *file, int in_place) {
	char command[8192];
	printf("Converting %s to baseline...",file);
	char *esc_file = escape(file);
	memset(command,0,sizeof(command));
	int result = snprintf(command,8192,"jpegtran \"%s\" > \"%s.new\"",esc_file,esc_file);
	if((result == -1) || (result == 8192)) return;
	if(system(command)) {
		printf("error!\n");
		#ifdef _WIN32
		memset(command,0,sizeof(command));
		sprintf(command,"del \"%s.new\"",esc_file);
		system(command);
		#else
		memset(command,0,sizeof(command));
		sprintf(command,"rm -f \"%s.new\"",esc_file);
		system(command);
		#endif
		return;
	}

	char *destination = NULL;
	if(in_place) {
		#ifdef _WIN32
		memset(command,0,sizeof(command));
		sprintf(command,"%s~",file);
		remove(command);
		char *drive = strrchr(file,'\\');
		if(drive) ++drive;
		else drive = file;
		memset(command,0,sizeof(command));
		sprintf(command,"rename \"%s\" \"%s~\"",esc_file,drive);
		system(command);
		#else
		memset(command,0,sizeof(command));
		sprintf(command,"mv -f \"%s\" \"%s~\"",esc_file,esc_file);
		system(command);
		#endif
	} else {
		char *new_filename = strrchr(esc_file,dir_sep);
		if(new_filename) ++new_filename;
		else new_filename = esc_file;

		destination = malloc(sizeof(char)*(strlen(new_filename)+8));
		sprintf(destination,"%s%c%s","output",dir_sep,new_filename);
	}
	memset(command,0,sizeof(command));
	#ifdef _WIN32
	if(!in_place) remove(destination);
	sprintf(command,"move \"%s.new\" \"%s\"",esc_file,(in_place?esc_file:destination));
	#else
	sprintf(command,"mv -f \"%s.new\" \"%s\"",esc_file,(in_place?esc_file:destination));
	#endif
	system(command);
	if(destination) free(destination);
	free(esc_file);
	printf("done!\n");
}

zipFile open_output_zip(char *file) {
	char *filext = strrchr(file,'.');
	int add = (filext?strlen(filext):4);
	if(add < 4) add = 4;
	char *new_file = malloc(sizeof(char)*(strlen(file)+strlen("output/")+add));
	sprintf(new_file,"output/%s",file);
	filext = strrchr(new_file,'.');
	if(!filext) {
		strcat(new_file,".zip");
	} else if(strcasecmp(filext,".zip") != 0) {
		filext[1] = 'z';
		filext[2] = 'i';
		filext[3] = 'p';
	}
	zipFile new_zip_file = zipOpen(new_file,APPEND_STATUS_CREATE);
	free(new_file);
	return new_zip_file;
}
void add_buffer_to_zip(zipFile zfile, char *in_arc, void *buffer, long file_size) {
	if(!(zfile && in_arc && buffer)) return;
	zipOpenNewFileInZip(zfile,in_arc,NULL,NULL,0,NULL,0,NULL,(level?Z_DEFLATED:0),level);
	zipWriteInFileInZip(zfile,buffer,file_size);
	zipCloseFileInZip(zfile);
}
void add_to_zip(zipFile zfile, char *on_disc, char *in_arc) {
	FILE *the_file = fopen(on_disc,"rb");
	if(!the_file) {
		printf("Error! Could not open file %s!\n",on_disc);
		return;
	}
	fseek(the_file,0,SEEK_END);
	long file_size = ftell(the_file);
	if(!file_size) {
		fclose(the_file);
		return;
	}
	void *buffer = malloc(file_size);
	fseek(the_file,0,SEEK_SET);
	fread(buffer,1,file_size,the_file);
	fclose(the_file);
	add_buffer_to_zip(zfile,in_arc,buffer,file_size);
}

inline void remove_backup(char *file) {
	if(!file) return;
	char *new_file = malloc(sizeof(char)*(strlen(file)+2));
	sprintf(new_file,"%s~",file);
	remove(new_file);
	free(new_file);
}

void run_zip(char *file) {
	printf("Processing zip file %s\n",file);
	voidp zip_file = unzOpen(file);
	if(!zip_file) {
		printf("Couldn't open zip file!\n");
		return;
	}
	zipFile new_zip_file = open_output_zip(strrchr(file,dir_sep)+1);
	if(!new_zip_file) {
		printf("Couldn't open output zip file!\n");
		unzClose(zip_file);
		return;
	}
	char more = 1;
	char internal_filename[4096];
	char *outfilename = NULL;
	while(more) {
		unz_file_info fi;
		memset(internal_filename,0,4096);
		if(unzGetCurrentFileInfo(zip_file,&fi,internal_filename,4096,NULL,0,NULL,0) < 0) {
			zipClose(new_zip_file,NULL);
			unzClose(zip_file);
			printf("Error handling zip file!\n");
			return;
		}
		printf("Extracting %s...",internal_filename);
		if(unzOpenCurrentFile(zip_file) < 0) {
			zipClose(new_zip_file,NULL);
			unzClose(zip_file);
			printf("Error handling zip file!\n");
			return;
		}
		voidp buffer = malloc(fi.uncompressed_size);
		if(!buffer) {
			zipClose(new_zip_file,NULL);
			unzCloseCurrentFile(zip_file);
			unzClose(zip_file);
			printf("Out of memory!\n");
			return;
		}
		unzReadCurrentFile(zip_file,buffer,fi.uncompressed_size);
		unzCloseCurrentFile(zip_file);
		outfilename = malloc(sizeof(char)*(strlen("temp")+strlen(internal_filename)+2));
		if(!outfilename) {
			free(buffer);
			zipClose(new_zip_file,NULL);
			unzClose(zip_file);
			printf("Out of memory!\n");
			break;
		}
		sprintf(outfilename,"temp%c%s",dir_sep,internal_filename);
		#ifdef _WIN32 //Needed on Windows but not on Unixes because backslashes are allowed in filenames
		char *working_filename = internal_filename;
		char *filext;
		while((filext = strchr(working_filename,'\\'))) {
			char *dirname = malloc(sizeof(char)*(strlen(outfilename)-strlen(filext)+1));
			if(!dirname) {
				zipClose(new_zip_file,NULL);
				unzClose(zip_file);
				printf("Out of memory!\n");
				return;
			}
			memcpy(dirname,outfilename,sizeof(char)*(strlen(outfilename)-strlen(filext)));
			dirname[strlen(outfilename)-strlen(filext)] = 0;
			if(mkdir(dirname) == -1) {
				if(errno != EEXIST) {
					free(dirname);
					zipClose(new_zip_file,NULL);
					unzClose(zip_file);
					printf("Could not create output directory!\n");
					return;
				}
			}
			free(dirname);
			working_filename = filext+1;
		}
		#endif
		FILE *outfile = fopen(outfilename,"wb");
		if(!outfile) {
			printf("Error opening output file!\n");
			printf("Let's just pretend this file is baseline and stick in back in anyway...");
			add_buffer_to_zip(new_zip_file, internal_filename, buffer,fi.uncompressed_size);
			free(buffer);
			printf("done!\n");
			if(unzGoToNextFile(zip_file) < 0) more = 0;
			continue;
		}
		fwrite(buffer,1,fi.uncompressed_size,outfile);
		fclose(outfile);
		free(buffer);
		printf("done!\n");
		if(unzGoToNextFile(zip_file) < 0) more = 0;
		if(is_jpeg(outfilename)) to_baseline(outfilename,1);
		printf("Adding file back to zip...");
		add_to_zip(new_zip_file,outfilename,internal_filename);
		printf("done!\n");
		remove(outfilename);
		remove_backup(outfilename);
	}
	zipClose(new_zip_file,NULL);
	unzClose(zip_file);
	printf("Done with zip file!\n");
}

void run_rar(char *file) {
	printf("Processing rar file %s\n",file);
	void *rar_file = open_rar(file);
	if(!rar_file) {
		printf("Couldn't open rar file!\n");
		return;
	}
	zipFile new_zip_file = open_output_zip(strrchr(file,dir_sep)+1);
	if(!new_zip_file) {
		printf("Couldn't open output zip file!\n");
		close_rar(rar_file);
		return;
	}
	char more = 1;
	char internal_filename[260];
	while(more) {
		unsigned int flags;
		int result = get_fileinfo_from_rar(rar_file,internal_filename,&flags);
		if(result > 10) { //10 signals end of archive, greater is an error
			zipClose(new_zip_file,NULL);
			close_rar(rar_file);
			printf("Error handling rar file!\n");
			return;
		} else if(result == 10) break;
		if(!((flags & RAR_FILE_IS_DIR) == RAR_FILE_IS_DIR)) {
			char *outfilename = malloc(sizeof(char)*(strlen("temp")+strlen(internal_filename)+2));
			sprintf(outfilename,"temp%c%s",dir_sep,internal_filename);
			printf("Extracting %s...",internal_filename);
			extract_file_from_rar(rar_file,"./",outfilename);
			printf("done!\n");
			if(is_jpeg(outfilename)) to_baseline(outfilename,1);
			printf("Adding file back to zip...");
			add_to_zip(new_zip_file,outfilename,internal_filename);
			printf("done!\n");
			remove(outfilename);
			remove_backup(outfilename);
			free(outfilename);
		}
	}
	zipClose(new_zip_file,NULL);
	close_rar(rar_file);
	printf("Done with rar file!\n");
}

void real_main(char *file, int in_place) {
	DIR *dirh = opendir(file);
	if(dirh) {
		struct dirent *de;
		while((de = readdir(dirh)) > 0) {
			size_t namelength;
			#ifdef _WIN32
			namelength = de->d_namlen;
			#else
			namelength = strlen(de->d_name);
			#endif
			char *filename = malloc(sizeof(char)*(namelength+1));
			strcpy(filename,de->d_name);
			if(strcmp(filename,".") && strcmp(filename,"..") &&
					(strcmp(file,".") || strcmp(filename,"temp") && strcmp(filename,"output"))) {
				if(filename[namelength-1] == dir_sep) filename[namelength-1] = 0;
				char *new_dir = malloc(sizeof(char)*(strlen(file)+strlen(filename)+2));
				sprintf(new_dir,"%s%c%s",file,dir_sep,filename);
				real_main(new_dir,in_place);
				free(new_dir);
			}
			free(filename);
		}
		closedir(dirh);
		return;
	}
	char *filext = strrchr(file,'.');
	if(!filext) return;
	if(!strcasecmp(filext,".jpg") || !strcasecmp(filext,".jpeg") || !strcasecmp(filext,".jpe"))
		to_baseline(file,in_place);
	else if(!strcasecmp(filext,".cbz") || !strcasecmp(filext,".zip"))
		run_zip(file);
	else if(!strcasecmp(filext,".cbr") || !strcasecmp(filext,".rar"))
		run_rar(file);
}

int main(int argc, char *argv[]) {
	//Detect if jpegtran is present
	const char *path = getenv("PATH");
	char *path_tok = malloc(sizeof(char)*(strlen(path)+1));
	if(!path_tok) return 1;
	sprintf(path_tok,"%s",path);
	char *token = strtok(path_tok,path_sep);
	int found = 0;
	while(token && !found) {
		char *filename = malloc(sizeof(char)*(strlen(token)+1+strlen("jpegtran")+strlen(exe_ext)+1));
		if(!filename) return 1;
		sprintf(filename,"%s%cjpegtran%s",token,dir_sep,exe_ext);
		struct stat s;
		if(stat(filename,&s) == 0) found = 1;
		free(filename);
		token = strtok(NULL,path_sep);
	}
	free(path_tok);
	if(!found) {
		printf("%s: jpegtran not found\n",argv[0]);
		return 1;
	}

	//Make temp and output directories
	#ifdef _WIN32
	if(mkdir("temp") < 0) {
	#else
	if(mkdir("temp",0777) < 0) {
	#endif
		if(errno != EEXIST) {
			printf("Could not create temp directory!\n");
			return 1;
		}
	}
	#ifdef _WIN32
	if(mkdir("output") < 0) {
	#else
	if(mkdir("output",0777) < 0) {
	#endif
		if(errno != EEXIST) {
			printf("Could not create output directory!\n");
			return 1;
		}
	}
	fclose(stdin);  //This is a hack to keep libunrar from clogging up the
	//fclose(stderr); //output and making the program interactive in the event
					// that the silent version of libunrar is not used
	if(argc > 1) {
		int i = 1, inp = 0;
		if(argv[1][0] == '-') {
			i = 2;
			if(strlen(argv[1]) < 2) {
				printf("%s: bad command line option -\n",argv[0]);
				return 2;
			}
			int j;
			for(j = 1; j < strlen(argv[1]); ++j) {
				if(isdigit(argv[1][j])) level = argv[1][1] - 0x30;
				else {
					switch(argv[1][j]) {
						case 'i': inp = 1; break;
						default:
							printf("%s: bad command line option %c\n",argv[0],argv[1][j]);
							return 2;
					}
				}
			}
		}
		if(i == argc) printf("%s: no input files specified\n",argv[0]);
		else for(; i < argc; ++i) real_main(argv[i],inp);
	} else printf("%s: no input files specified\n",argv[0]);

	rmdir("temp");
	return 0;
}
