// affinity.cpp : Defines the entry point for the application.
//

#include "stdafx.h"
#include "affinity.h"
#include "shellbrowse.h"

INT_PTR CALLBACK AffinityDlgProc(HWND, UINT, WPARAM, LPARAM);

int APIENTRY _tWinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPTSTR lpCmdLine, int nCmdShow)
{
	DialogBox(hInstance, MAKEINTRESOURCE(IDD_AFFINITY), NULL, AffinityDlgProc);
	return 0;
}

tstring GetLastErrorString(DWORD dwLastError)
{
	TCHAR* buffer = 0;

	::FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
		NULL, dwLastError, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (PTSTR)&buffer, 0, NULL);

	tstring retval;

	if (buffer != 0)
	{
		retval = buffer;
		LocalFree(buffer);
	}

	return retval;
}

void DisplayErrorMessage(HWND hWnd, const TCHAR* message, DWORD dwLastError)
{
	tstring caption = message;

	caption += _T(": ");
	caption += GetLastErrorString(dwLastError);

	MessageBox(hWnd, caption.c_str(), _T("Error"), MB_OK | MB_ICONERROR);
}

int CALLBACK BrowseCallbackProc(HWND hwnd, UINT uMsg, LPARAM lParam, LPARAM lpData)
{
	HWND hWndDialog = (HWND)lpData;

	TCHAR szDir[MAX_PATH] = { 0 };
	GetDlgItemText(hWndDialog, IDC_PATH, szDir, MAX_PATH);

    if (uMsg == BFFM_INITIALIZED)
		SendMessage(hwnd, BFFM_SETSELECTION, TRUE, (LPARAM)szDir);

	return 0;
}

bool Affinity_OnCommand(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam, LRESULT& lResult)
{
	bool bHandled = true;

	switch (LOWORD(wParam))
	{
	case IDC_PATH_BROWSE:
		{
			TCHAR szFilePath[MAX_PATH] = { 0 };

			OPENFILENAME ofn = { 0 };
			ofn.lStructSize = sizeof(ofn);
			ofn.hwndOwner = hDlg;
			ofn.lpstrFile = szFilePath;
			ofn.nMaxFile = MAX_PATH;
			ofn.lpstrFilter = _T("Executables (*.exe)\0*.EXE\0All Files (*.*)\0*.*\0");
			ofn.nFilterIndex = 0;

			TCHAR szCurrentDirectory[MAX_PATH] = { 0 };
			GetCurrentDirectory(MAX_PATH, szCurrentDirectory);
			ofn.lpstrInitialDir = szCurrentDirectory;

			ofn.Flags = OFN_PATHMUSTEXIST | OFN_FILEMUSTEXIST;

			if (GetOpenFileName(&ofn))
				SetDlgItemText(hDlg, IDC_PATH, szFilePath);

			break;
		}
	case IDC_WORKDIR_BROWSE:
		{
			BROWSEINFO bi = { 0 };

			bi.ulFlags = BIF_USENEWUI|BIF_VALIDATE;
			bi.lpszTitle = _T("Choose a working directory:");
			bi.lpfn = BrowseCallbackProc;
			bi.lParam = (LPARAM)hDlg;

			LPITEMIDLIST pIDL = SHBrowseForFolder(&bi);

			if (pIDL != 0)
			{
				TCHAR szWorkDir[MAX_PATH] = { 0 };

				if (SUCCEEDED(SHGetTargetFolderPath(pIDL, szWorkDir, MAX_PATH)))
					SetDlgItemText(hDlg, IDC_WORKDIR, szWorkDir);

				CoTaskMemFree(pIDL);
			}

			break;
		}
	case IDC_FIRSTONLY:
		for (UINT id = IDC_CHECK2; id <= IDC_CHECK32; id++)
			CheckDlgButton(hDlg, id, BST_UNCHECKED);
		break;
	case IDC_ALLCORES:
		{
			for (UINT id = IDC_CHECK1; id <= IDC_CHECK32; id++)
			{
				if (IsWindowEnabled(GetDlgItem(hDlg, id)))
					CheckDlgButton(hDlg, id, BST_CHECKED);
			}
		}
		break;
	case IDCANCEL:
		EndDialog(hDlg, IDCANCEL);
		break;
	case IDOK:
		{
			DWORD dwAffinity = 0;

			for (UINT id = IDC_CHECK1; id <= IDC_CHECK32; id++)
			{
				if (IsWindowEnabled(GetDlgItem(hDlg, id)))
				{
					if (IsDlgButtonChecked(hDlg, id) == BST_CHECKED)
						dwAffinity |= (1 << (id - IDC_CHECK1));
				}
			}

			if (dwAffinity == 0)
				dwAffinity = 1;

			TCHAR szPath[MAX_PATH] = { 0 };
			TCHAR szArguments[MAX_PATH] = { 0 };
			TCHAR szWorkingDir[MAX_PATH] = { 0 };

			GetDlgItemText(hDlg, IDC_PATH, szPath, MAX_PATH);
			GetDlgItemText(hDlg, IDC_ARGUMENTS, szArguments, MAX_PATH);
			GetDlgItemText(hDlg, IDC_WORKDIR, szWorkingDir, MAX_PATH);

			if (!_tcscmp(szWorkingDir, _T("")))
			{
				// if there is no working dir, get the dir of the program.
				_tcsncpy(szWorkingDir, szPath, MAX_PATH);
				szWorkingDir[MAX_PATH - 1] = 0;

				// if there is no dir specified for the program, get this process's working dir
				if (!PathRemoveFileSpec(szWorkingDir))
					GetCurrentDirectory(MAX_PATH, szWorkingDir);
			}

			STARTUPINFO info = { 0 };
			PROCESS_INFORMATION psinfo = { 0 };

			// lol pyramid code
			if (CreateProcess(szPath, szArguments, NULL, NULL, FALSE, CREATE_SUSPENDED, NULL, szWorkingDir, &info, &psinfo))
			{
				if (SetProcessAffinityMask(psinfo.hProcess, dwAffinity))
				{
					if (ResumeThread(psinfo.hThread) == 0xFFFFFFFF)
					{
						DisplayErrorMessage(hDlg, _T("ResumeThread failed"), GetLastError());
					}
					else
					{
						EndDialog(hDlg, LOWORD(wParam));
					}
				}
				else
				{
					DisplayErrorMessage(hDlg, _T("SetProcessAffinityMask failed"), GetLastError());
				}
			}
			else
			{
				DisplayErrorMessage(hDlg, _T("CreateProcess failed"), GetLastError());
			}
		}
		break;
	default:
		bHandled = false;
		break;
	}

	lResult = (INT_PTR)TRUE;
	return bHandled;
}

// Message handler for about box.
INT_PTR CALLBACK AffinityDlgProc(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam)
{
	UNREFERENCED_PARAMETER(lParam);
	switch (message)
	{
	case WM_INITDIALOG:
		{
			DWORD_PTR dwProcessAffinity = 0;
			DWORD_PTR dwSystemAffinity = 0;
			if (!GetProcessAffinityMask(GetCurrentProcess(), &dwProcessAffinity, &dwSystemAffinity))
			{
				MessageBox(hDlg, _T("Could not get system affinity mask."), _T("Error"), MB_OK | MB_ICONERROR);
				EndDialog(hDlg, IDCANCEL);
			}

			for (UINT id = IDC_CHECK1; id <= IDC_CHECK32; id++)
			{
				if (!(dwSystemAffinity & (1 << (id - IDC_CHECK1))))
					EnableWindow(GetDlgItem(hDlg, id), FALSE);

				TCHAR buffer[8] = { 0 };
				wsprintf(buffer, _T("CPU %d"), id - IDC_CHECK1);

				SetDlgItemText(hDlg, id, buffer);
			}

			return (INT_PTR)TRUE;
		}
	case WM_COMMAND:
		{
			LRESULT lResult = 0;
			if (Affinity_OnCommand(hDlg, message, wParam, lParam, lResult))
				return lResult;
		}
		break;
	}

	return (INT_PTR)FALSE;
}