#include "pch.h"
#include "ProcessHelper.h"
#include <windows.h>
#include <tlhelp32.h>
#include <tchar.h>
#include <strsafe.h>

// 请求关机或重启，EWX_SHUTDOWN、EWX_REBOOT
// cmd: EWX_SHUTDOWN, EWX_REBOOT, EWX_POWEROFF 等
// reason: SHTDN_REASON_MAJOR_OPERATINGSYSTEM, SHTDN_REASON_MINOR_OTHER 等
void DoShutdown(uint32_t cmd, DWORD reason)
{
	HANDLE hToken;
	TOKEN_PRIVILEGES tkp;

	// 获取进程令牌
	if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
		_tprintf(_T("OpenProcessToken failed (%d)\n"), GetLastError());
		return;
	}

	// 获取关机权限
	LookupPrivilegeValue(NULL, SE_SHUTDOWN_NAME, &tkp.Privileges[0].Luid);
	tkp.PrivilegeCount = 1;
	tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

	// 调整权限
	if (!AdjustTokenPrivileges(hToken, FALSE, &tkp, 0, NULL, 0)) {
		_tprintf(_T("AdjustTokenPrivileges failed (%d)\n"), GetLastError());
		CloseHandle(hToken);
		return;
	}

	// 关闭令牌句柄
	CloseHandle(hToken);

	// 调用关机或重启
	if (!ExitWindowsEx(cmd, reason)) {
		_tprintf(_T("ExitWindowsEx failed (%d)\n"), GetLastError());
	}
}



// 检查某个进程是否正在运行，返回找到的第一个匹配进程名
BOOL IsProcessRunning(LPCTSTR pszProcName, LPTSTR pFoundName, DWORD dwFoundLen)
{
	HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hSnapshot == INVALID_HANDLE_VALUE)
		return FALSE;

	PROCESSENTRY32 pe;
	pe.dwSize = sizeof(PROCESSENTRY32);

	if (Process32First(hSnapshot, &pe))
	{
		do
		{
			if (_tcsicmp(pe.szExeFile, pszProcName) == 0)
			{
				StringCchCopy(pFoundName, dwFoundLen, pe.szExeFile);
				CloseHandle(hSnapshot);
				return TRUE;
			}
		} while (Process32Next(hSnapshot, &pe));
	}

	CloseHandle(hSnapshot);
	return FALSE;
}

// 查找正在运行的目标进程，返回数量，并将找到的进程写入 outFoundProcs
int FindRunningProcesses_BeginWith(const std::vector<std::wstring>& targetProcs, std::vector<sProcessInfo>& outFoundProcs)
{
	// outFoundProcs大小要么与targetProcs相同，要么等于空
	bool sameSize = (targetProcs.size() == outFoundProcs.size());
	if (!sameSize && !outFoundProcs.empty()) {
		return -1; // 如果不相同且outFoundProcs不为空，则返回0
	}
	int count = 0;
	HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hSnapshot == INVALID_HANDLE_VALUE)
		return -2;

	PROCESSENTRY32W pe; // 使用宽字符版本
	pe.dwSize = sizeof(PROCESSENTRY32W);

	if (Process32FirstW(hSnapshot, &pe)) // 宽字符 API
	{
		do
		{
			std::wstring currProcName(pe.szExeFile);

			for (uint32_t i = 0; i < targetProcs.size(); i++)
			{
				const auto& target = targetProcs[i];
				if (currProcName.size() >= target.size() && 0 == _wcsnicmp(currProcName.c_str(), target.c_str(), target.size()))
				{
					++count;
					if (sameSize) {
						outFoundProcs[i] = sProcessInfo(currProcName, pe.th32ProcessID); // 更新对应位置
					}
					else {
						outFoundProcs.push_back(sProcessInfo(currProcName, pe.th32ProcessID));
					}

					break;
				}
			}

		} while (Process32NextW(hSnapshot, &pe));
	}

	CloseHandle(hSnapshot);
	return count;
}

// 查找正在运行的目标进程，返回数量，并将找到的进程写入 outFoundProcs
int FindRunningProcesses_Exact(const std::vector<std::wstring>& targetProcs, std::vector<sProcessInfo>& outFoundProcs)
{
	// outFoundProcs大小要么与targetProcs相同，要么等于空
	bool sameSize = (targetProcs.size() == outFoundProcs.size());
	if (!sameSize && !outFoundProcs.empty()) {
		return -1; // 如果不相同且outFoundProcs不为空，则返回0
	}
	int count = 0;
	HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hSnapshot == INVALID_HANDLE_VALUE)
		return -2;

	PROCESSENTRY32W pe; // 使用宽字符版本
	pe.dwSize = sizeof(PROCESSENTRY32W);

	if (Process32FirstW(hSnapshot, &pe)) // 宽字符 API
	{
		do
		{
			std::wstring currProcName(pe.szExeFile);

			for (uint32_t i = 0; i < targetProcs.size(); i++)
			{
				const auto& target = targetProcs[i];
				if (0 == _wcsicmp(currProcName.c_str(), target.c_str()))
				{
					++count;
					if (sameSize) {
						outFoundProcs[i] = sProcessInfo(currProcName, pe.th32ProcessID); // 更新对应位置
					}
					else {
						outFoundProcs.push_back(sProcessInfo(currProcName, pe.th32ProcessID));
					}

					break;
				}
			}

		} while (Process32NextW(hSnapshot, &pe));
	}

	CloseHandle(hSnapshot);
	return count;
}

// 启动一个进程，并获取其 PID。注意启动的程序是当前程序的子进程。
bool StartProcessAndGetPid(const fs::path& exePath, DWORD* outPid)
{
	// 将路径转换为宽字符串
	std::wstring widePath = exePath.wstring();

	STARTUPINFO si = { sizeof(STARTUPINFO) };
	PROCESS_INFORMATION pi;

	// 创建进程
	if (CreateProcess(
		NULL,                   // 应用程序名称（NULL 表示使用命令行）
		&widePath[0],           // 命令行（可修改参数）
		NULL,                   // 进程句柄不可继承
		NULL,                   // 线程句柄不可继承
		FALSE,                  // 不继承句柄
		0,                      // 无创建标志
		NULL,                   // 使用父进程环境
		NULL,                   // 使用父进程目录
		&si,                    // 启动信息
		&pi                     // 返回的进程信息
	))
	{
		// 获取 PID
		if (outPid) *outPid = pi.dwProcessId;

		// 关闭进程和线程句柄（但不等待进程结束）
		CloseHandle(pi.hProcess);
		CloseHandle(pi.hThread);

		return true;
	}
	else
	{
		return false;
	}
}

// 独立启动一个进程，使其与当前进程解耦（类似双击运行）
bool StartProcessIndependently(const fs::path& exePath)
{
	// 构造命令：cmd /c start "" "full\path\to\a.exe"
	// 注意：start 的第一个参数是窗口标题，留空用 ""
	std::wstring cmd = L"cmd.exe /c start \"\" \"" + exePath.wstring() + L"\"";

	STARTUPINFOW si = { sizeof(STARTUPINFOW) };
	PROCESS_INFORMATION pi = {};

	// 启动 cmd.exe（隐藏窗口，不继承句柄）
	BOOL success = CreateProcessW(
		nullptr,                // lpApplicationName
		&cmd[0],                // lpCommandLine (可修改)
		nullptr,                // lpProcessAttributes
		nullptr,                // lpThreadAttributes
		FALSE,                  // bInheritHandles = false
		CREATE_NO_WINDOW |      // 隐藏中间 cmd.exe 的黑窗口（用户体验更好）
		DETACHED_PROCESS,       // 进一步解耦确保新进程不继承调用者的控制台（即使 cmd 本身有控制台也不传给 a.exe）
		nullptr,                // lpEnvironment
		nullptr,                // lpCurrentDirectory
		&si,                    // lpStartupInfo
		&pi                     // lpProcessInformation
	);

	if (success)
	{
		CloseHandle(pi.hProcess);
		CloseHandle(pi.hThread);
		return true;
	}

	return false;
}

bool NotifyProcessToExit(DWORD pid, DWORD timeoutMs)
{
	HWND mainHwnd = nullptr;
	struct FindMainWindowParam {
		DWORD pid;
		HWND  hwnd;
	} param{ pid, nullptr };

	EnumWindows([](HWND hwnd, LPARAM lParam) -> BOOL {
		auto* p = reinterpret_cast<FindMainWindowParam*>(lParam);
		DWORD processId = 0;
		GetWindowThreadProcessId(hwnd, &processId);
		if (processId == p->pid) {
			if (GetWindow(hwnd, GW_OWNER) == nullptr && IsWindowVisible(hwnd)) {
				p->hwnd = hwnd;
				return FALSE; // 找到主窗口，停止枚举
			}
		}
		return TRUE;
		}, reinterpret_cast<LPARAM>(&param));

	mainHwnd = param.hwnd;

	if (mainHwnd) {
		PostMessage(mainHwnd, WM_CLOSE, 0, 0);
	}
	else {
		// 枚举该进程的所有窗口并发送 WM_CLOSE
		EnumWindows([](HWND hwnd, LPARAM lParam) -> BOOL {
			DWORD processId = 0;
			GetWindowThreadProcessId(hwnd, &processId);
			if (processId == static_cast<DWORD>(lParam)) {
				PostMessage(hwnd, WM_CLOSE, 0, 0);
			}
			return TRUE;
			}, static_cast<LPARAM>(pid));
	}

	// 等待进程退出
	HANDLE hProcess = OpenProcess(SYNCHRONIZE, FALSE, pid);
	if (hProcess)
	{
		DWORD result = WaitForSingleObject(hProcess, timeoutMs);
		CloseHandle(hProcess);
		return result == WAIT_OBJECT_0;
	}

	return false;
}