#include <stdio.h>
#include <io.h>
#include <stdlib.h>
#include <windows.h>
#include "uplaunch.h"
extern "C" {
#include "zlib.h"
}
//Python vars & functions
DECLVAR (int, Py_NoSiteFlag);
DECLVAR (int, Py_VerboseFlag);
DECLPROC (int, Py_Initialize, (void));
DECLPROC (int, Py_Finalize, (void));
DECLPROC (PyObject *, Py_CompileString, (char *, char *, int));
DECLPROC (PyObject *, PyImport_ExecCodeModule, (char *, PyObject *));
DECLPROC (int, PyRun_SimpleString, (char *));
DECLPROC (int, PySys_SetArgv, (int, char **));
DECLPROC (void, Py_SetProgramName, (char *));
DECLPROC (PyObject *, PyImport_ImportModule, (char *));
DECLPROC (int, PyObject_SetAttrString, (PyObject *, char *, PyObject *));
DECLPROC (PyObject *, PyList_New, (int));
DECLPROC (int, PyList_Append, (PyObject *, PyObject *));
DECLPROC (PyObject *, Py_BuildValue, (char *, ...));
DECLPROC (PyObject *, PyFile_FromFile, (FILE *, char *, char *, int));
DECLPROC (PyObject *, PyObject_CallFunction, (PyObject *, char *, ...));
DECLPROC (PyObject *, PyModule_GetDict, (PyObject *));
DECLPROC (PyObject *, PyDict_GetItemString, (PyObject *, char *));
DECLPROC (PyObject *, PyErr_Occurred, () );
DECLPROC (void, PyErr_Print, () );
DECLPROC (void, Py_SetPythonHome, (char *));
DECLPROC (char *, Py_GetPath, (void));

//Forward
int getSelf(void);
int openSelf(void);
int runIt(int argc, char *argv[]);
HINSTANCE loadPython(void);
unsigned char *extract(struct TOC *toc, FILE *fp, int pkgstart);
#ifndef _WRAP_PYTHON_
int extract2fs(struct TOC *toc, FILE *fp, int pkgstart, char *homepath, char *temppath);
void cleanUp();
char temppath[_MAX_PATH+1];
#endif
void import(struct TOC *toc, FILE *fp, int pkgstart);
int install(int pkgstart, char* path, char *arname);
int run(struct TOC *toc, FILE *fp, int pkgstart);
typedef int (__cdecl *PYMAIN)(int, char **);
//Globals
char thisfile[_MAX_PATH+1];
FILE *fp;			/* this file */
int pkgstart;
char *TOCbuff = NULL;
char homepath[_MAX_PATH+1];
char *workpath = NULL;		/* ends up pointing to homepath or temppath */
struct COOKIE c;
HINSTANCE dll;

#ifdef _CONSOLE
int main(int argc, char* argv[])
#else
int APIENTRY WinMain( HINSTANCE hInstance, HINSTANCE hPrevInstance,
						LPSTR lpCmdLine, int nCmdShow )
#endif
{
	int rc = 0;
	VS("Start " STUBNAME "\n");
	/* this fills in thisfile, homepath and temppath */
	if (getSelf()) return -1;

	/* this sets fp, fills in pkgstart, COOKIE c, and TOCBuff */
	VS("Opening self\n");
	if (openSelf()) return -1;

#ifndef _WRAP_PYTHON_
	int dorun = 1;
	int doextract = 1;
	workpath = getenv( "_MEIPASS2" );
	if (workpath)
		doextract = 0;

	//pass 1 - unpack all the binary members (pythonxx.dll, zlib.pyd etc...)
	if (doextract) {
		VS("Unpacking binaries\n");
		struct TOC *toc = (struct TOC *)TOCbuff;
		while ((char*)toc < TOCbuff + c.TOClen) {
			if (toc->typcd == 'b') {
				dorun = 0;
				if (!extract2fs(toc, fp, pkgstart, homepath, temppath)) 
					FATALERROR("Could not extract binaries.\n");
			}
			toc = (struct TOC *)((char*)toc + toc->structlen);
			if (toc == NULL) {
				FATALERROR("Cannot read Table of Contents.\n");
				return -1;
			}
		}
	}

	if (dorun) {
#endif
		/* we didn't extract anything */
#ifndef _CONSOLE
		rc = runIt(__argc, __argv);
#else
		rc = runIt(argc, argv);
#endif
#ifndef _WRAP_PYTHON_
	}
	else {
		SECURITY_ATTRIBUTES sa;
		sa.nLength = sizeof(sa);
		sa.lpSecurityDescriptor = NULL;
		sa.bInheritHandle = TRUE;
		STARTUPINFO si;
		GetStartupInfo(&si);
		si.lpReserved = NULL;
		si.lpDesktop = NULL;
		si.lpTitle = NULL;
		si.dwFlags = STARTF_USESTDHANDLES | STARTF_USESHOWWINDOW;
		si.wShowWindow = SW_NORMAL;
		si.hStdInput = (void*)_get_osfhandle(fileno(stdin));
		si.hStdOutput = (void*)_get_osfhandle(fileno(stdout));
		si.hStdError = (void*)_get_osfhandle(fileno(stderr));
		/* tell pass 2 where we extracted to */
		char *envvar = (char *)malloc(strlen(workpath)+10);
		strcpy(envvar, "_MEIPASS2=");
		strcat(envvar, workpath);
		_putenv(envvar);
		PROCESS_INFORMATION pi;
		if (CreateProcess( 
				thisfile, // pointer to name of executable module 
				GetCommandLine(),  // pointer to command line string 
				&sa,  // pointer to process security attributes 
				NULL,  // pointer to thread security attributes 
				TRUE,  // handle inheritance flag 
				0,  // creation flags 
				NULL,  // pointer to new environment block 
				NULL,  // pointer to current directory name 
				&si,  // pointer to STARTUPINFO 
				&pi  // pointer to PROCESS_INFORMATION 
				)) {
			WaitForSingleObject(pi.hProcess, INFINITE);
			GetExitCodeProcess(pi.hProcess, (unsigned long *)&rc);
		}
		else {
			FATALERROR("Error creating child process!\n");
			rc = -1;
		}
		free(envvar);
	}
	
	if (doextract && !dorun) {
		VS("Removing binaries\n");
		cleanUp();
	}
#endif
	return rc;
}

int runIt(int argc, char *argv[])
{
	int	rc = 0, i;
	char *pypath = NULL;	
	dll = loadPython();
	if (dll) {
	//now launch python ...
#ifndef _WRAP_PYTHON_
	if (workpath) {
		/* we extracted something */
		if (stricmp(workpath, temppath)==0) {
			/* need both - may have .pyd's on homepath */
			pypath = (char*)malloc(strlen(temppath)+strlen(homepath)+13);
			strcpy(pypath, "PYTHONPATH=");
			strcat(pypath, temppath);
			strcat(pypath, ";");
			strcat(pypath, homepath);
		}
		else {
			/* workpath is homepath - no need for temppath */
			pypath = (char*)malloc(strlen(homepath)+13);		
			strcpy(pypath, "PYTHONPATH=");
			strcat(pypath, homepath);
		}
	}
	else {
		/* never extracted anything - homepath will d	o */
		pypath = (char*)malloc(strlen(homepath)+13);
		strcpy(pypath, "PYTHONPATH=");
		strcat(pypath, homepath);
	}
	pypath[strlen(pypath)-1] = '\0';
	VS("Manipulating evironment\n");
	_putenv(pypath);
	_putenv("PYTHONHOME=");
#else	
	pypath = (char*)malloc(strlen(homepath)+13);
	strcpy(pypath,"PYTHONHOME=");
	strcpy(pypath+11,homepath);
	i = strlen(pypath)-1;
	if(pypath[i]=='\\' || pypath[i]=='/') pypath[i]=0;
	VS1("Set PYTHONHOME=%s\n",pypath);
	_putenv( pypath );
#endif	
	if(pypath) free(pypath);

		VS("Setting Python vars\n");
#if !defined(_WRAP_PYTHON_)
		*Py_NoSiteFlag = 1;
#else
		*Py_NoSiteFlag = 0;
#endif

		int vflag_index = -1;
		VS1("Py_NoSiteFlag=%d\n", *Py_NoSiteFlag);
		for (i = 0; i < argc; ++i) {
			if (0 == strcmp (argv[i], "-v")) {
				*Py_VerboseFlag = 1;	
				vflag_index = i;
			}
		}
		
		Py_SetProgramName(thisfile);
		VS("Initializing Python\n");	
		Py_Initialize();
    	Py_GetPath();
    	_putenv( "PYTHONPATH=" );	
#if !defined(_WRAP_PYTHON_)
		VS("Manipulating Python's sys.path\n");
		PyRun_SimpleString("import sys;del sys.path[1:]");
#else	
		VS("Importing module sys\n");
		PyRun_SimpleString("import sys");
#endif

		PyObject *py_argv = PyList_New (0);
		PyObject *val = Py_BuildValue ("s", thisfile);
		PyList_Append (py_argv, val);
		for (i = 1; i < argc; ++i) {
			if (i != vflag_index) {	
				PyObject *val = Py_BuildValue ("s", argv[i]);
				PyList_Append (py_argv, val);
			}
		}
		PyObject *sys = PyImport_ImportModule ("sys");
		VS("Setting sys.argv\n");
		PyObject_SetAttrString (sys, "argv", py_argv);
		// pass 2 - "import" all the modules
		VS("importing modules from CArchive\n");
		PyObject *marshal = PyImport_ImportModule("marshal");
		PyObject *marshaldict = PyModule_GetDict(marshal);
		PyObject *loadfunc = PyDict_GetItemString(marshaldict, "load");
		PyObject *pyfile = PyFile_FromFile(fp, thisfile, "rb", 0);
		struct TOC *toc = (struct TOC *)TOCbuff;
		char *q;
		char buff[40];
		char cmd[80];
		while ((char*)toc < TOCbuff + c.TOClen) {	
			if (toc->typcd == 'm' || toc->typcd == 'M') {
#ifdef _DEBUG
				fprintf(stderr, "loading %s\n", toc->name);
#endif
				fseek(fp, pkgstart + toc->pos + 8, SEEK_SET);
				PyObject *co = PyObject_CallFunction(loadfunc, "O", pyfile);
				PyObject *mod = PyImport_ExecCodeModule (toc->name, co);
				if( mod == NULL )
					fprintf(stderr, "mod is NULL loading %s\n", toc->name );
				if( PyErr_Occurred() )
					{
					fprintf(stderr, "error occurred loading %s\n", toc->name );
					PyErr_Print();
					}

				if (q=strrchr(toc->name, '.')) {
					strncpy(buff, toc->name, q - toc->name);
					buff[q - toc->name] = '\0';	
					sprintf(cmd, "sys.modules[\"%s\"]%s = sys.modules[\"%s\"]", buff, q, toc->name);
#ifdef _DEBUG
					fprintf(stderr, "Run_SimpleString( %s )\n", cmd);
#endif
					PyRun_SimpleString(cmd);
				}
			}
			toc = (struct TOC *)((char*)toc + toc->structlen);
			if (toc == NULL) {
				FATALERROR("Cannot read Table of Contents.\n");
				return -1;
			}
		}
		// pass 3 - install all the zlibs
		VS("Installing import hooks\n");
		toc = (struct TOC *)TOCbuff;
		while ((char*)toc < TOCbuff + c.TOClen) {
			if (toc->typcd == 'z'){
				VS2("install(%d,,%s)\n",toc->pos+pkgstart, thisfile);
				install(toc->pos+pkgstart, "", thisfile);
				}
			toc = (struct TOC *)((char*)toc + toc->structlen);
			if (toc == NULL) {
				FATALERROR("Cannot read Table of Contents.\n");
				return -1;
			}
		}
		// pass 4 - run any scripts
		VS("Running scripts\n");
		toc = (struct TOC *)TOCbuff;
		while ((char*)toc < TOCbuff + c.TOClen) {
			if (toc->typcd == 's') 
				run(toc, fp, pkgstart);
			toc = (struct TOC *)((char*)toc + toc->structlen);
			if (toc == NULL) {
				FATALERROR("Cannot read Table of Contents.\n");
				return -1;
			}
		}
		VS("Finalizing Python\n");
		Py_Finalize ();
	}


	VS("closing self\n");
	fclose(fp);
	VS1("Returning %d to OS\n", rc);
	return rc;
}
//forward
unsigned char *decompress(unsigned char *, struct TOC *);
	
int getSelf(void) 
{
	if (!GetModuleFileNameA(NULL, thisfile, _MAX_PATH)) {
		FATALERROR("System error - unable to load!");
		return -1;
	}
	char *p = thisfile+strlen(thisfile) - 4;
	if (strnicmp(p, ".exe", 4) != 0)
		strcat(thisfile, ".exe");
#ifndef	_WRAP_PYTHON_
	//GetModuleFileName returns an absolute path
	strcpy(homepath, thisfile);
	for (p=homepath+strlen(homepath); *p != '\\' && p >= homepath; --p);
	*++p = '\0';

	//Get path to official temp directory
	GetTempPath(MAX_PATH, temppath);
#endif

	return 0;
}

int openSelf(void)
{
	fp = fopen(thisfile, "rb");
	if (fp == NULL) {
		FATALERROR("Cannot open self: ");
		FATALERROR(thisfile);
		return -1;
	}
	fseek(fp, 0, SEEK_END);
	int filelen = ftell(fp);
	if (fseek(fp, -(int)sizeof(COOKIE), SEEK_END)) {
		FATALERROR("Archive appears to be invalid");
		return -1;
	}
	fread(&c, sizeof(COOKIE), 1, fp);
	if (strncmp(c.magic, MAGIC, 8)) {
		FATALERROR("Archive has bad magic - quiting!");
		return -1;
	}
	pkgstart = filelen - c.len;
	fseek(fp, pkgstart+c.TOC, SEEK_SET);
	TOCbuff = (char *)malloc(c.TOClen);
	if (TOCbuff == NULL) {
		FATALERROR("Could not allocate buffer for TableOfContents\n");
		return -1;
	}
	fread(TOCbuff, c.TOClen, 1, fp);
	struct TOC *toc = (struct TOC *)TOCbuff;

#ifdef	_WRAP_PYTHON_
	{
	long	rc, bufSize=_MAX_PATH;
	char	*subKeyFMT = "Software\\Python\\PythonCore\\%d.%d\\InstallPath";
	char	*subKey = "Software\\Python\\PythonCore\\  .  \\InstallPath";
	sprintf(subKey,subKeyFMT,c.pyvers/10, c.pyvers % 10);
	VS1("Finding RegKey HKEY_LOCAL_MACHINE\\%s\n",subKey);
	if((rc = RegQueryValue(HKEY_LOCAL_MACHINE, subKey, homepath, &bufSize))!=ERROR_SUCCESS){
		FATALERROR("Can't Find PYTHONHOME\n");
		return -1;
		}
	}
#endif
	return 0;
}

HINSTANCE loadPython(void) {
	HINSTANCE dll;
#ifdef _WRAP_PYTHON_
	char *dllpath = (char*)malloc(strlen(homepath)+20);
#else
	char *dllpath = (char*)malloc(max(strlen(homepath), strlen(temppath))+20);
#endif
	VS("Loading python\n");
	sprintf(dllpath, "%spython%02d.dll", homepath, c.pyvers);
	VS(dllpath);VS("\n");
	dll = LoadLibraryEx(dllpath, NULL, LOAD_WITH_ALTERED_SEARCH_PATH);
#ifndef _WRAP_PYTHON_
	if (!dll) {
		sprintf(dllpath, "%spython%02d.dll", temppath, c.pyvers);
		VS(dllpath);
		VS("\n");
		dll = LoadLibraryEx(dllpath, NULL, LOAD_WITH_ALTERED_SEARCH_PATH);
		}
#else
	if (!dll) {
		VS("Attempting global Python\n");
		sprintf(dllpath, "python%02d.dll", c.pyvers);
		dll = LoadLibrary(dllpath);
		}
#endif
	if(!dll){
		VS("Failed to load python dll\n");
		}
	else{
		// Thomas Heller's code
		VS("Getting entry points\n");
		GETVAR (dll, int, Py_NoSiteFlag);
		GETVAR (dll, int, Py_VerboseFlag);
		GETPROC (dll, int, Py_Initialize, (void));
		GETPROC (dll, int, Py_Finalize, (void));	
		GETPROC (dll, PyObject *, Py_CompileString, (char *, char *, int));
		GETPROC (dll, PyObject *, PyImport_ExecCodeModule, (char *, PyObject *));
		GETPROC (dll, int, PyRun_SimpleString, (char *));
		GETPROC (dll, int, PySys_SetArgv, (int, char **));
		GETPROC (dll, void, Py_SetProgramName, (char *));
		GETPROC (dll, PyObject *, PyImport_ImportModule, (char *));
		GETPROC (dll, int, PyObject_SetAttrString, (PyObject *, char *, PyObject *));
		GETPROC (dll, PyObject *, PyList_New, (int));
		GETPROC (dll, int, PyList_Append, (PyObject *, PyObject *));
		GETPROC (dll, PyObject *, Py_BuildValue, (char *, ...));
		GETPROC (dll, PyObject *, PyFile_FromFile, (FILE *, char *, char *, int));
		GETPROC (dll, PyObject *, PyObject_CallFunction, (PyObject *, char *, ...));
		GETPROC (dll, PyObject *, PyModule_GetDict, (PyObject *));
		GETPROC (dll, PyObject *, PyDict_GetItemString, (PyObject *, char *));
		GETPROC (dll, PyObject *, PyErr_Occurred, () );
		GETPROC (dll, void, PyErr_Print, () );
		GETPROC (dll, void, Py_SetPythonHome, (char *));
		GETPROC (dll, char *, Py_GetPath, (void));
		}
	free(dllpath);

	return dll;
}

int install(int pkgstart, char* path, char *arname)
{
	int r;
	char *tmpl = "import archive_rt;sys.importers.append(archive_rt.ZlibArchive(r'%s%s', %d))\n";
	char *cmd = (char *)malloc(strlen(tmpl) + strlen(arname) + strlen(path) + 32);
	sprintf(cmd, tmpl, path, arname, pkgstart);
#if _DEBUG
	fprintf(stderr, "executing %s\n", cmd);
#endif
	r = PyRun_SimpleString(cmd);
	if(r){
		cmd = (char*)realloc(cmd,strlen(cmd)+25);
		strcat(cmd," !!!FAILED!!!\n");
		VS(cmd);
		}
	else {
		VS(cmd);
		}
	free(cmd);
	return r;
}

int run(struct TOC *toc, FILE *fp, int pkgstart)
{
	int r;
	unsigned char *data = extract(toc, fp, pkgstart);
#ifdef _DEBUG
	fprintf(stderr, "running %s\n", toc->name);
#endif
	r = PyRun_SimpleString((char *)data);
#ifdef _DEBUG
	fprintf(stderr, "%s\n", r ? "FAILED" : "SUCCEEDED" );
#endif
	free(data);
	return r;
}
unsigned char *extract(struct TOC *toc, FILE *fp, int pkgstart)
{
#if _DEBUG
	fprintf(stderr, " extracting %s (%d, %c)\n", toc->name, toc->cflag, toc->typcd);
#endif
	fseek(fp, pkgstart + toc->pos, SEEK_SET);
	unsigned char *data = (unsigned char *)malloc(toc->len);
	if (data == NULL) {
		OTHERERROR("Could not allocate read buffer\n");
		return NULL;
	}
	fread(data, toc->len, 1, fp);
	if (toc->cflag == '\1') {
		unsigned char *tmp = decompress(data, toc);
		free(data);
		data = tmp;
		if (data == NULL) {
			char msg[40];
			sprintf(msg, "Error decompressing %s\n", toc->name);
			OTHERERROR(msg);
			return NULL;
		}
	}
	return data;
}

unsigned char *decompress(unsigned char * buff, struct TOC *toc)
{
	const char *ver = (zlibVersion)();
	unsigned char *out = (unsigned char *)malloc(toc->ulen);
	if (out == NULL) {
		OTHERERROR("Error allocating decompression buffer\n");
		return NULL;
	}
	z_stream zstream;

	zstream.zalloc = NULL;
	zstream.zfree = NULL;
	zstream.opaque = NULL;
	zstream.next_in = buff;
	zstream.avail_in = toc->len;
	zstream.next_out = out;
	zstream.avail_out = toc->ulen;
	int rc = inflateInit(&zstream);
	if (rc >= 0) {
		rc = (inflate)(&zstream, Z_FINISH);
		if (rc >= 0) {
			rc = (inflateEnd)(&zstream);
		}
		else {
			char msg[40];
			sprintf(msg, "Error %d from inflate: %s\n", rc, zstream.msg);
			OTHERERROR(msg);
			return NULL;
		}
	}
	else {
		char msg[40];
		sprintf(msg, "Error %d from inflateInit: %s\n", rc, zstream.msg);
		OTHERERROR(msg);
		return NULL;
	}	
	return out;
}

#ifndef _WRAP_PYTHON_
static	int removeFromFS(struct TOC *toc,const char* rootPath)
{
	int	r;
	char name[_MAX_PATH+1];
	strcpy(name, rootPath);
	strcat(name, toc->name);
	HMODULE h = GetModuleHandle(name);
	if(h){
		int i=0;
		while(FreeLibrary(h) && i<100) i++;
		VS2("Library %s freed %d times\n",name,i); 
		}
	r = remove(name);
	VS2("File '%s' removed\n",name, r ? "not" : "");
	return r;
}

void cleanUp()
{
	struct TOC *toc = (struct TOC *)TOCbuff;
	if (toc == NULL) {
		VS("Cannot read Table of Contents!\n");
		}
	else{
		while((char*)toc < TOCbuff + c.TOClen){
			if (toc->typcd == 'b'){
				if (removeFromFS(toc,homepath)==-1){
					removeFromFS(toc,temppath);
					}
				}
			toc = (struct TOC *)((char *)toc + toc->structlen);
			if (toc == NULL) {
				VS("exitFunc: Table of Contents went bad!\n");
				}
			}
		}
	int tries = 20;
	int rc = 1;
	while (tries-- && rc) rc = FreeLibrary(dll);
  	//free(TOCbuff);
}

int extract2fs(struct TOC *toc, FILE *fp, int pkgstart, char *homepath, char *temppath)
{
	char fnm[_MAX_PATH+1];
	FILE *out;
	unsigned char *data = extract(toc, fp, pkgstart);
	if (!workpath) {
		strcpy(fnm, homepath);
		strcat(fnm, toc->name);
		out = fopen(fnm, "wb");
		if (out == NULL) {
			strcpy(fnm, temppath);
			strcat(fnm, toc->name);
			out = fopen(fnm, "wb");
			if (out == NULL) {
				free(data);
				return 0;
			}
			workpath = temppath;
		}
		else
			workpath = homepath;
	}
	else {
		strcpy(fnm, workpath);
		strcat(fnm, toc->name);
		out = fopen(fnm, "wb");
	}
	fwrite(data, toc->ulen, 1, out);
	fclose(out);
	free(data);
	return 1;
}
#endif
