#include "precompile.h"
#include "osutil.h"


#ifdef _WIN32

#include <TlHelp32.h>

TOOLKIT_API int osutil_detect_unique_app(char** pNames, int nNum)
{
	//const WCHAR *exe;
	HANDLE hSnapshot;
	int rc = TRUE;

	/*{
		WCHAR *t = (WCHAR*)_alloca(MAX_PATH);
		GetModuleFileNameW(NULL, t, MAX_PATH);
		exe = wcsrchr(t, '\\') + 1;
	}*/
	DWORD dwCurProcID = GetCurrentProcessId();

	hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hSnapshot) {
		PROCESSENTRY32 pe;
		pe.dwSize = sizeof(pe);
		if (Process32First(hSnapshot, &pe)) {
			do {
				int i;
				for (i = 0; i < nNum; i++) {
					if (stricmp(&pe.szExeFile[0], pNames[i]) == 0 && pe.th32ProcessID != dwCurProcID) {
						rc = FALSE;
						break;
					}
				}
			} while (Process32Next(hSnapshot, &pe));
		}
		CloseHandle(hSnapshot);
	}

	return rc;
}

TOOLKIT_API int osutil_restart_system()
{
	HANDLE hToken;
	TOKEN_PRIVILEGES tkp;
	BOOL fResult;

	// Get the current process token handle so we can get shutdown privilege. 
	if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
		fprintf(stderr, "get proc token fail");
		return -1;
	}
	// Get the LUID for shutdown privilege. 
	LookupPrivilegeValue(NULL, SE_SHUTDOWN_NAME, &tkp.Privileges[0].Luid);

	tkp.PrivilegeCount = 1;  // one privilege to set    
	tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

	// Get shutdown privilege for this process. 
	AdjustTokenPrivileges(hToken, FALSE, &tkp, 0, (PTOKEN_PRIVILEGES)NULL, 0);
	// Cannot test the return value of AdjustTokenPrivileges. 
	if (GetLastError() != ERROR_SUCCESS) {
		fprintf(stderr, "adjust proc token privilege fail");
		return -1;
	}

	// Display the shutdown dialog box and start the countdown. 
	fResult = InitiateSystemShutdown(
		NULL,    // shut down local computer 
		NULL,   // message for user
		0,			// time-out period, in seconds 
		FALSE,   // ask user to close apps 
		TRUE);   // reboot after shutdown 

	if (!fResult) {
		fprintf(stderr, "request windows reboot fail");
		return -1;
	}
	// Disable shutdown privilege. 
	tkp.Privileges[0].Attributes = 0;
	AdjustTokenPrivileges(hToken, FALSE, &tkp, 0, (PTOKEN_PRIVILEGES)NULL, 0);
	return 0;
}

TOOLKIT_API void osutil_terminate_related_process(char** process_array, const int array_size)
{
	const DWORD dwCurProcessID = GetCurrentProcessId();
	DWORD relateProcessIDs[256];
	DWORD dwIdx = 0;
	char szCmd[256];
	DWORD relateProcessNum = 0;

	HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

	memset(relateProcessIDs, 0, sizeof(relateProcessIDs));

	if (hSnapshot)
	{
		PROCESSENTRY32 pe;
		pe.dwSize = sizeof(pe);

		if (Process32First(hSnapshot, &pe)) {
			do {
				int i;
				for (i = 0; i < array_size; i++) {
					if (stricmp(&pe.szExeFile[0], process_array[i]) == 0 && pe.th32ProcessID != dwCurProcessID) {
						relateProcessIDs[relateProcessNum++] = pe.th32ProcessID;
						break;
					}
				}
			} while (Process32Next(hSnapshot, &pe));
		}
		CloseHandle(hSnapshot);
	}
	
	for (dwIdx = 0; dwIdx < relateProcessNum; ++dwIdx) {
		sprintf_s(szCmd, 256, "TASKKILL /PID %d /F", relateProcessIDs[dwIdx]);
		WinExec(szCmd, SW_HIDE);
	}
}

#else

#include <dirent.h>
#include <unistd.h>
#include <strings.h>
#include <winpr/wtypes.h>

static inline int is_num(const char* chars)
{
	for (; *chars; chars++)
		if (*chars < '0' || *chars > '9')
			return 0; // false
	return 1; // true
}

TOOLKIT_API int osutil_detect_unique_app(char** pNames, int nNum)
{
	int rc = TRUE;
	DIR* dir_proc = NULL;
	struct dirent* dir_entity = NULL;
	int i = 0;
	char cmdLinePath[512] = {0};
	char processName[128] = {0};
	char processPath[256] = {0};
	pid_t  pid = (pid_t)-1;
	const pid_t cur_pid = getpid();
	dir_proc = opendir("/proc/");
	if (dir_proc == NULL) {
		return FALSE;
	}
	while ((dir_entity = readdir(dir_proc))) {
		if (dir_entity->d_type == DT_DIR) {
			if (is_num(dir_entity->d_name)) {
				pid = (pid_t)atoi(dir_entity->d_name);
				if (pid == getpid()) {
					continue;
				}
				strcpy(cmdLinePath, "/proc/");
				strcat(cmdLinePath, dir_entity->d_name);
				strcat(cmdLinePath, "/cmdline");
				FILE* fd_cmdline = fopen(cmdLinePath, "rt");
				if (fd_cmdline) {
					fscanf(fd_cmdline, "%s", processPath);
					fclose(fd_cmdline);
					if (strrchr(processPath, '/')) {
						strcpy(processName, strrchr(processPath, '/') + 1);
					}
					else {
						strcpy(processName, processPath);
					}
					for (i = 0; i < nNum; i++) {
						if (strcasecmp(processName, pNames[i]) == 0) {
							printf("%s, %s, %d, %d\n", processPath, processName, pid, cur_pid);
							rc = FALSE;
							break;
						}
					}
				}
			}
		}
		if (!rc) {
			break;
		}
	}
	closedir(dir_proc);
	return rc;
}

TOOLKIT_API int osutil_restart_system()
{
	//TODO: implement it.
	return 0;
}

TOOLKIT_API void osutil_terminate_related_process(char** process_array, const int array_size)
{
	
}

#endif