#include <pspkernel.h>
#include <pspdebug.h>
#include <pspiofilemgr.h>
#include <psputils.h>
#include <pspctrl.h>
#include <psppower.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

PSP_MODULE_INFO("piKey Installer", 0x1000, 1, 1);
PSP_MAIN_THREAD_ATTR(THREAD_ATTR_USER);

#define	SIZEDIGEST	16

#define DATA_MAIN   0

typedef struct {
	char *src;
	char *dest;
	char *MD5;
} FILE_DATA;

FILE_DATA data[] = {
	{ "Modules/piKey.PRX",					"MS0:/SEPlugins/piKey.PRX",									"4e347a96f2ddd86a7d577881cb80c833"},
	{ "Modules/piKeyConfig.txt",			"MS0:/SEPlugins/piKey/piKeyconfig.txt",						"6f04bd1f90ec6436272fde809b265dd1"},
	{ "Modules/CTRLMap.txt",				"MS0:/SEPlugins/piKey/CTRLMap.txt",							"e2b97ec1cf2d87387201a76ce5b06048"},
	{ "Modules/piKeyCTRLOut.PRX",			"MS0:/SEPlugins/piKey/OutputDrivers/piKeyCTRLOut.PRX",		"876e8bcc6172342ad5043defecc0d584"},
	{ "Modules/piKeyDanzeffOut.PRX",		"MS0:/SEPlugins/piKey/OutputDrivers/piKeyDanzeffOut.PRX",	"859539b5e903898d143adf434add1cb7"},
	{ "Modules/piKeyOSKOut.PRX",			"MS0:/SEPlugins/piKey/OutputDrivers/piKeyOSKOut.PRX",		"e3f8f8c70892e7ba29c32692f641436d"},
	{ "Modules/piKeySIO.PRX",				"MS0:/SEPlugins/piKey/InputDrivers/piKeySIO.PRX",			"bdd7c5509d0e4d2b6ce08de5104d8196"},
	{ "Modules/piKeyIR.PRX",				"MS0:/SEPlugins/piKey/InputDrivers/piKeyIR.PRX",			"6420f2c8fbf6cc84c07ee4ffe2b2d16d"},
	{ "Modules/PSPIRKeyB.ini",				"MS0:/SEPlugins/piKey/PSPIRKeyB.ini",						"5fe651cf91c998855bb6c8295b611bd8"},
	{ "Modules/Keymap/Targus.ini",			"MS0:/SEPlugins/piKey/Keymap/Targus.ini",					"9c8743af38d7d11598dc2c1d68b8e416"},
	{ "Modules/Keymap/Qwerty.ini",		 	"MS0:/SEPlugins/piKey/Keymap/Qwerty.ini",					"e5dd92c99e65cf4cc49cd7d2e3c6111d"},
	{ "Modules/Keymap/Default.ini",			"MS0:/SEPlugins/piKey/Keymap/Default.ini",					"11c9b068161b3f7219695e3ac648b280"},
	{ "Modules/Keymap/Hama.ini",			"MS0:/SEPlugins/piKey/Keymap/Hama.ini",						"f1b010c117dbc38dd90edece80d8a05f"},
	{ "Modules/Keymap/Novaets.ini",			"MS0:/SEPlugins/piKey/Keymap/Novaets.ini",					"9790ed397c1010c010b95c894aead91e"},
	{ "Modules/Keymap/Palm.ini",			"MS0:/SEPlugins/piKey/Keymap/Palm.ini",						"9790ed397c1010c010b95c894aead91e"},
	{ "Modules/Keymap/Keymap Readme.txt",	"MS0:/SEPlugins/piKey/Keymap/Keymap Readme.txt",			"9790ed397c1010c010b95c894aead91e"},
	{ NULL, NULL, NULL }
};

/*****************************************************************************/
/* append text to a file.                                                    */
/*****************************************************************************/
void append_txt(char *filename, char *txt)
{
	int fd = sceIoOpen(filename,
								PSP_O_RDONLY | PSP_O_WRONLY | PSP_O_APPEND | PSP_O_CREAT,
								0777);

	sceIoLseek32(fd, 0, SEEK_END);

	sceIoWrite(fd, txt, strlen(txt));

	sceIoClose(fd);
}

/*****************************************************************************/
/* returns 0 for success.                                                    */
/*****************************************************************************/
int check_md5(char *file, char *md5){
 	u8 digest[SIZEDIGEST];
	int i, fd, size;
	char *buffer;
	char digestBuff[(SIZEDIGEST*2) + 1];
	
	memset(digest, 0, SIZEDIGEST);
	fd = sceIoOpen(file, PSP_O_RDONLY, 0);
	if (fd < 0) {
		pspDebugScreenPrintf("Cannot open PRX.");
		return 1;
	}
		
	size = sceIoLseek32(fd, 0, PSP_SEEK_END);
	sceIoLseek32(fd, 0, PSP_SEEK_SET);
	size -= 16;
		
	buffer = malloc(size);
		
	if(!buffer) {
		sceIoClose(fd);
		pspDebugScreenPrintf("Error allocating buffer.");
		return 1;
	}
		
	sceIoRead(fd, buffer, size);
	sceIoClose(fd);
	
	sceKernelUtilsMd5Digest((u8 *) buffer, size, digest);

	free(buffer);
	
	for(i = 0; i < 16; i++)
		sprintf(&digestBuff[i * 2], "%02x", (int) digest[i]);
	
	if(strncmp(md5, digestBuff, SIZEDIGEST*2) != 0)
	{
		append_txt("MS0:/md5.txt", "\nFailed MD5: ");
		append_txt("MS0:/md5.txt", file);
		append_txt("MS0:/md5.txt", "\nmd5a: ");
		append_txt("MS0:/md5.txt", md5);
		append_txt("MS0:/md5.txt", "\nmd5b: ");
		append_txt("MS0:/md5.txt", digestBuff);
		return 1;
	}
	
	return 0;
}

void create_dirs(void) {
	SceUID fd;
	
	fd = sceIoDopen("MS0:/SEPlugins");
	if(fd < 0) sceIoMkdir("MS0:/SEPlugins", 0777);
	sceIoDclose(fd);
	
	fd = sceIoDopen("MS0:/SEPlugins/piKey");
	if(fd < 0) sceIoMkdir("MS0:/SEPlugins/piKey", 0777);
	sceIoDclose(fd);
	
	fd = sceIoDopen("MS0:/SEPlugins/piKeyInputDriver");
	if(fd < 0) sceIoMkdir("MS0:/SEPlugins/piKey/InputDrivers", 0777);
	sceIoDclose(fd);
	
	fd = sceIoDopen("MS0:/SEPlugins/piKey/OutputDrivers");
	if(fd < 0) sceIoMkdir("MS0:/SEPlugins/piKey/OutputDrivers", 0777);
	sceIoDclose(fd);

	fd = sceIoDopen("MS0:/SEPlugins/piKey/Keymap");
	if(fd < 0) sceIoMkdir("MS0:/SEPlugins/piKey/Keymap", 0777);
	sceIoDclose(fd);
}
	
	

void copy_file(char *src, char *dst) {
	SceUID infd = sceIoOpen(src, PSP_O_RDONLY, 0777);
	char buffer[8192];
	int bytesread, totalwritten = 0;
	sceIoRemove(dst);
	SceUID outfd = sceIoOpen(dst, PSP_O_WRONLY | PSP_O_CREAT | PSP_O_TRUNC, 0777);
	
	if (infd < 0 || outfd < 0) return;
	
	while ((bytesread = sceIoRead(infd, buffer, 8192)) > 0)
		totalwritten += sceIoWrite(outfd, buffer, bytesread);
	
	sceIoClose(infd);
	sceIoClose(outfd);
}

void checkPower(void) {
	if(!scePowerIsPowerOnline()) {
		pspDebugScreenPrintf("Please connect the AC adapter and restart.\n"
				    "Exiting in 10 seconds.");
		
		sceKernelDelayThread(10 * 1000 * 1000);
		sceKernelExitGame();
	}
}

int flash_piKey(void) {
	int rc, failed = 0;

	rc = sceIoUnassign("flash0:");
	if(rc < 0)
		failed = 1;
	
	if (!failed)
	{
		rc  = sceIoAssign("flash0:", "lflash0:0,0", "flashfat0:", 0, IOASSIGN_RDWR , 0);
		if(rc < 0)
			failed = 1;
	}
	
	if(!failed)
		copy_file(data[DATA_MAIN].src, "flash0:/kd/piKey.PRX");
	else
		return 1;
	
	return 0;
}

int main(void) {
	SceCtrlData pad;
	int flashed = 0, i, fd;
	
	pspDebugScreenInit();
	
	pspDebugScreenPrintf("piKey Installer v0.4\n"
					"Welcome to the 'Unofficial piKey 0.4 Installer.'\n\n"
				"Press 'X' to install.\n"
"Press 'Home' to exit.\n\n");

for(;;)
	{
		sceCtrlReadBufferPositive(&pad, 1);
		if(pad.Buttons & PSP_CTRL_HOME)
		{
			pspDebugScreenPrintf("Installation Canceled\n"
			     "Exiting in 10 seconds.");
	
	sceKernelDelayThread(10 * 1000 * 1000);
	sceKernelExitGame();
	return 0;
}
	else if(pad.Buttons & PSP_CTRL_CROSS)
		{
			pspDebugScreenPrintf("Continuing Installation\n\n");
			break;
		}
	}

	pspDebugScreenPrintf("Creating directories... \n");
	
	create_dirs();
	pspDebugScreenPrintf("Done!\n"
			     "Copying files... ");
	for(i = 0; data[i].src != NULL; i++)
		copy_file(data[i].src, data[i].dest);
	
	pspDebugScreenPrintf("Done!\n");

#if 1
	pspDebugScreenPrintf("\n"
			     "Do you want to install piKey.PRX for Game150 homebrew?\n"
				 "(It's only required if your PSP kernel is 3.40 or below. \n"
				 "(Requires writing to flash0) We will not be held responsible for any unlikely bricks.\n"
			     "Press 'X' for yes or 'O' for no.\n");
	
	for(;;)
	{
		sceCtrlReadBufferPositive(&pad, 1);
		if(pad.Buttons & PSP_CTRL_CROSS)
		{
			pspDebugScreenPrintf("Copying to flash0... ");
			if(flash_piKey() > 0)
			{
				pspDebugScreenPrintf("Failed!\n\n");
				sceKernelDelayThread(10 * 1000 * 1000);
				sceKernelExitGame();
			}
			else
			{
				// and now write to pspbtcnf.txt

				fd = sceIoOpen("flash0:/kd/pspbtcnf_Game.txt",
											PSP_O_RDONLY | PSP_O_WRONLY | PSP_O_APPEND | PSP_O_CREAT,
											0777);

				// Read file, check for existing piKey entry, copying text to the new file as we go
				int needToAppend = 1;
				int lsize = sceIoLseek32(fd, 0, SEEK_END);
				if (fd >= 0)
				{
					char *lfilebuff = malloc(lsize + 1);
					sceIoLseek32(fd, 0, SEEK_SET);
					sceIoRead(fd, lfilebuff, lsize);
					lfilebuff[lsize] = '\0';

					if (strstr(lfilebuff, "piKey.PRX"))
					{
						// no need to append config, it's already there.
						pspDebugScreenPrintf("Done!\n\n");
						needToAppend = 0;
					}

					free(lfilebuff);
				}
				else
				{
					pspDebugScreenPrintf("Failed to open pspbtcnf_Game.txt: %08X\n", fd);
				}

				if (needToAppend)
				{
					// Need to insert our text after #/kd/systemctrl.PRX, which is near
					// the end of the file.
#define FINDSTR "#/kd/systemctrl.PRX\r\n"

					char lbuff[512];
					int lpos = lsize - 512;
					int found = 0;

					while (!found)
					{
						sceIoLseek32(fd, lpos, SEEK_SET);
						if (sceIoRead(fd, lbuff, sizeof(FINDSTR)) == 0)
						{
							// EOF
							pspDebugScreenPrintf("Failed to find the marker text in pspbtcnf_Game.txt!\n");
							found = 1;
							lpos = -1;
						}
						else
						{
							if (memcmp(lbuff, FINDSTR, sizeof(FINDSTR) - 1) == 0)
							{
								// found the right offset. seek forwards to the next byte.
								sceIoLseek32(fd, lpos + sizeof(FINDSTR) - 1, SEEK_SET);
								found = 1;
							}
							else
							{
								lpos++;
							}
						}
					}

					if (lpos != -1)
					{
						// we found the text, read and save the remainder of the file
						lsize = sceIoRead(fd, lbuff, 512);
						sceIoLseek32(fd, lpos + sizeof(FINDSTR) - 1, SEEK_SET);
						sceIoWrite(fd, "/kd/piKey.PRX\r\n", sizeof("/kd/piKey.PRX\r\n"));
						sceIoWrite(fd, lbuff, lsize);
						flashed = 1;
						pspDebugScreenPrintf("Done!\n\n");
					}
				}

				sceIoClose(fd);

				break;
			}
		}
		else if(pad.Buttons & PSP_CTRL_CIRCLE)
		{
			pspDebugScreenPrintf("Not enabling for Game150\n\n");
			break;
		}
	}
#endif
	
	pspDebugScreenPrintf("Adding to MS0:/SEPlugins/VSH.txt... ");
	
	fd = sceIoOpen("MS0:/SEPlugins/VSH.txt",
								PSP_O_RDONLY | PSP_O_WRONLY | PSP_O_APPEND | PSP_O_CREAT,
								0777);

	// Read file, check for existing piKey entry, copying text to the new file as we go
	int needToAppend = 1;
	if (fd >= 0)
	{
		int lsize = sceIoLseek32(fd, 0, SEEK_END);
		char *lfilebuff = malloc(lsize + 1);
    sceIoLseek32(fd, 0, SEEK_SET);
		sceIoRead(fd, lfilebuff, lsize);
		lfilebuff[lsize] = '\0';

		if (strstr(lfilebuff, "piKey.PRX"))
		{
			// no need to append config, it's already there.
			needToAppend = 0;
		}

		free(lfilebuff);
	}
	else
	{
		pspDebugScreenPrintf("Failed to open VSH.txt: %08X\n", fd);
	}

	if (needToAppend)
	{
		sceIoLseek32(fd, 0, SEEK_END);
  	sceIoWrite(fd, "\nMS0:/SEPlugins/piKey.PRX\n", sizeof("\nMS0:/SEPlugins/piKey.PRX\n") - 1);
	}
	
	sceIoClose(fd);
	
	pspDebugScreenPrintf("Done!\n"
			     "Adding to MS0:/SEPlugins/Game.txt... ");
	
	fd = sceIoOpen("MS0:/SEPlugins/Game.txt",
								PSP_O_RDONLY | PSP_O_WRONLY | PSP_O_APPEND | PSP_O_CREAT,
								0777);

	// Read file, check for existing piKey entry, copying text to the new file as we go
	needToAppend = 1;
	if (fd >= 0)
	{
		int lsize = sceIoLseek32(fd, 0, SEEK_END);
		char *lfilebuff = malloc(lsize + 1);
    sceIoLseek32(fd, 0, SEEK_SET);
		sceIoRead(fd, lfilebuff, lsize);
		lfilebuff[lsize] = '\0';

		if (strstr(lfilebuff, "piKey.PRX"))
		{
			// no need to append config, it's already there.
			needToAppend = 0;
		}

		free(lfilebuff);
	}
	else
	{
		pspDebugScreenPrintf("Failed to open Game.txt: %08X\n", fd);
	}

	if (needToAppend)
	{
		sceIoLseek32(fd, 0, SEEK_END);
  	sceIoWrite(fd, "\nMS0:/SEPlugins/piKey.PRX\n", sizeof("\nMS0:/SEPlugins/piKey.PRX\n") - 1);
	}
	
	sceIoClose(fd);
	
	pspDebugScreenPrintf("Done!\n"
			     "Adding to MS0:/SEPlugins/Game150.txt... ");
	
	fd = sceIoOpen("MS0:/SEPlugins/Game150.txt",
								PSP_O_RDONLY | PSP_O_WRONLY | PSP_O_APPEND | PSP_O_CREAT,
								0777);

	// Read file, check for existing piKey entry, copying text to the new file as we go
	needToAppend = 1;
	if (fd >= 0)
	{
		int lsize = sceIoLseek32(fd, 0, SEEK_END);
		char *lfilebuff = malloc(lsize + 1);
    sceIoLseek32(fd, 0, SEEK_SET);
		sceIoRead(fd, lfilebuff, lsize);
		lfilebuff[lsize] = '\0';

		if (strstr(lfilebuff, "piKey.PRX"))
		{
			// no need to append config, it's already there.
			needToAppend = 0;
		}

		free(lfilebuff);
	}
	else
	{
		pspDebugScreenPrintf("Failed to open Game150.txt: %08X\n", fd);
	}

	if (needToAppend)
	{
		sceIoLseek32(fd, 0, SEEK_END);
  	sceIoWrite(fd, "\nMS0:/SEPlugins/piKey.PRX\n", sizeof("\nMS0:/SEPlugins/piKey.PRX\n") - 1);
	}

	sceIoClose(fd);

	pspDebugScreenPrintf("Done!\n\n"
			     "Thank you for installing 'Unofficial piKey 0.4'. Please enable piKey in the recovery menu!\n\n"
			     "Exiting in 10 seconds.");
	
	sceKernelDelayThread(10 * 1000 * 1000);
	sceKernelExitGame();
	return 0;
}
