#undef _UNICODE

#pragma comment (lib, "tier0.lib")
#pragma comment (lib, "tier1.lib")
#pragma comment (lib, "vstdlib.lib")
#pragma comment (linker, "/NODEFAULTLIB:libcmt")

#define CLIENT_DLL
//#define GMOD_BETA

//String stuff
#include <string>
#include <sstream>

//Lua module interface
#include "gmod/GMLuaModule.h"
#include <windows.h>
#include <TlHelp32.h>
#include <psapi.h>
#include <tchar.h>
#include <iostream>
#include <cstdlib>
#include "dbg.h"
#include "color.h"
#define PROCESS_METAID 2355
GMOD_MODULE( Open, Close );
LUA_FUNCTION(process_getall){
	BOOL            working = 0;
	PROCESSENTRY32  lppe = {0};
	DWORD           targetPid = 0;
	HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if(hSnapshot) {
		ILuaObject* ret=Lua()->GetNewTable();
		lppe.dwSize = sizeof(lppe);
		working=Process32First(hSnapshot, &lppe);
		int i=0;
		while(working) {
			i++;
			ILuaObject* proc=Lua()->GetNewTable();
			proc->SetMember("exefile",lppe.szExeFile);
			proc->SetMember("usage",(float)lppe.cntUsage);
			proc->SetMember("threads",(float)lppe.cntThreads);
			proc->SetMember("size",(float)lppe.dwSize);
			proc->SetMember("pid",(float)lppe.th32ProcessID);
			proc->SetMember("parentpid",(float)lppe.th32ParentProcessID);
			proc->SetMember("threads",(float)lppe.cntThreads);
			proc->SetMember("moduleid",(float)lppe.th32ModuleID);
			proc->SetMember("flags",(float)lppe.dwFlags);
			proc->SetMember("heapid",(float)lppe.th32DefaultHeapID);
			proc->SetMember("threadpriority",(float)lppe.pcPriClassBase);
			ret->SetMember(i,proc);
			working = Process32Next(hSnapshot, &lppe);
		}
		Lua()->Push(ret);
		CloseHandle(hSnapshot);
		return 1;
	}else{
		return 0;
	}
	return 0;
}

LUA_FUNCTION(process_getbypid){
	int pid=Lua()->GetInteger(1);
	HANDLE proc=OpenProcess(PROCESS_ALL_ACCESS, TRUE, pid);
	if(proc){
		ILuaObject* PROCESS_META=Lua()->GetMetaTable("Process", PROCESS_METAID);
		Lua()->PushUserData(PROCESS_META, proc);
		PROCESS_META->UnReference();
		return 1;
	}
	return 0;
}
LUA_FUNCTION(process_kill){
	Lua()->CheckType(1, PROCESS_METAID);
	int exitcode=Lua()->GetType(2)==GLua::TYPE_NUMBER ? Lua()->GetNumber(2) : 0;
	HANDLE* handle=static_cast<HANDLE*>(Lua()->GetUserData(1));
	Lua()->Push(TerminateProcess(handle, exitcode)!=0);
	return 1;
}
LUA_FUNCTION(process_getpath){
	Lua()->CheckType(1, PROCESS_METAID);
	int exitcode=Lua()->GetType(2)==GLua::TYPE_NUMBER ? Lua()->GetNumber(2) : 0;
	HANDLE* handle=static_cast<HANDLE*>(Lua()->GetUserData(1));
	LPSTR Buffer;
	if(GetModuleFileNameEx(handle, 0, Buffer, MAX_PATH)){
		Lua()->Push(reinterpret_cast<char*>(Buffer));
		return 1;
	}
	return 0;
}
LUA_FUNCTION(process_getmeminfo){
	Lua()->CheckType(1, PROCESS_METAID);
	int exitcode=Lua()->GetType(2)==GLua::TYPE_NUMBER ? Lua()->GetNumber(2) : 0;
	HANDLE* handle=static_cast<HANDLE*>(Lua()->GetUserData(1));
	PROCESS_MEMORY_COUNTERS pmc;
	if(GetProcessMemoryInfo(handle, &pmc, sizeof(pmc))){
		ILuaObject* ret=Lua()->GetNewTable();
		ret->SetMember("PageFaultCount",(float)pmc.PageFaultCount);
		ret->SetMember("PeakWorkingSetSize",(float)pmc.PeakWorkingSetSize);
		ret->SetMember("WorkingSetSize",(float)pmc.WorkingSetSize);
		ret->SetMember("QuotaNonPagedPoolUsage",(float)pmc.QuotaNonPagedPoolUsage);
		ret->SetMember("QuotaPeakPagedPoolUsage",(float)pmc.QuotaPeakPagedPoolUsage);
		ret->SetMember("QuotaPagedPoolUsage",(float)pmc.QuotaPagedPoolUsage);
		ret->SetMember("QuotaPeakNonPagedPoolUsage",(float)pmc.QuotaPeakNonPagedPoolUsage);
		ret->SetMember("PagefileUsage",(float)pmc.PagefileUsage);
		ret->SetMember("PeakPagefileUsage",(float)pmc.PeakPagefileUsage);
		Lua()->Push(ret);
		return 1;
	}
	return 0;
}
LUA_FUNCTION(process_getmodinfo){
	Lua()->CheckType(1, PROCESS_METAID);
	int exitcode=Lua()->GetType(2)==GLua::TYPE_NUMBER ? Lua()->GetNumber(2) : 0;
	HANDLE* handle=static_cast<HANDLE*>(Lua()->GetUserData(1));
	HMODULE modules[1024];
	DWORD cbNeeded;
	if(EnumProcessModules(handle, modules, sizeof(modules), &cbNeeded)){
		ILuaObject* ret=Lua()->GetNewTable();
		for ( int i = 0; i < (cbNeeded / sizeof(HMODULE)); i++ ){
			TCHAR Buffer[MAX_PATH];
			if(GetModuleFileNameEx(handle, modules[i], Buffer, sizeof(Buffer)/sizeof(TCHAR))){
				ret->SetMember(i+1,reinterpret_cast<char*>(Buffer));
			}
		}
		Lua()->Push(ret);
		return 1;
	}
	return 0;
}
LUA_FUNCTION(process_getname){
	Lua()->CheckType(1, PROCESS_METAID);
	int exitcode=Lua()->GetType(2)==GLua::TYPE_NUMBER ? Lua()->GetNumber(2) : 0;
	HANDLE* handle=static_cast<HANDLE*>(Lua()->GetUserData(1));
	LPSTR Buffer;
	if(GetModuleBaseName(handle, 0, Buffer, MAX_PATH)){
		Lua()->Push(reinterpret_cast<char*>(Buffer));
		return 1;
	}
	return 0;
}
LUA_FUNCTION(process_gc){
	Lua()->CheckType(1, PROCESS_METAID);
	int exitcode=Lua()->GetType(2)==GLua::TYPE_NUMBER ? Lua()->GetNumber(2) : 0;
	HANDLE* handle=static_cast<HANDLE*>(Lua()->GetUserData(1));
	CloseHandle(handle);
	return 0;
}
int Open(lua_State *L) {
	
	Color Blue( 0, 162, 232, 255 );
	Color White ( 255, 255, 255, 255 );
	HANDLE hToken;
	LUID sedebugnameValue;
	TOKEN_PRIVILEGES tkp;
	
	if(!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
		ConColorMsg(Blue, "gm_process");
		ConColorMsg( White, ": Couldn't get process token! Not running as administrator?\n" );
		return 0;
	}
	
	if(!LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &sedebugnameValue)) {
		ConColorMsg(Blue, "gm_process");
		ConColorMsg( White, ": Couldn't get privilege value! Not running as administrator?\n" );
		CloseHandle(hToken);
		return 0;
	}
	tkp.PrivilegeCount = 1;
	tkp.Privileges[0].Luid = sedebugnameValue;
	tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
	if(!AdjustTokenPrivileges(hToken, FALSE, &tkp, sizeof tkp, NULL, NULL)) {
		ConColorMsg(Blue, "gm_process");
		ConColorMsg( White, ": Couldn't adjust token privileges! Not running as administrator?\n" );
		CloseHandle(hToken);
		return 0;
	}
	Lua()->NewGlobalTable("process");
	ILuaObject* obj=Lua()->GetGlobal("process");
		obj->SetMember("GetAll",process_getall);
		obj->SetMember("GetByPID",process_getbypid);
		
	ILuaObject* PROCESS_META=Lua()->GetMetaTable("Process",PROCESS_METAID);
		PROCESS_META->SetMember("Kill",process_kill);
		PROCESS_META->SetMember("GetPath",process_getpath);
		PROCESS_META->SetMember("GetName",process_getname);
		PROCESS_META->SetMember("GetMemoryInfo",process_getmeminfo);
		PROCESS_META->SetMember("GetModuleInfo",process_getmodinfo);
		PROCESS_META->SetMember("Close",process_gc);
		PROCESS_META->SetMember("__gc",process_gc);
		
		PROCESS_META->SetMember("__index",PROCESS_META);
	ConColorMsg( Blue, "gm_process" );
	ConColorMsg( White, ": Loaded!\n" );

	return 0;
}

int Close(lua_State *L) {
	return 0;
}