#include "stdafx.h"
#include "ProcessAndModule.h"
#include "OSVersion.h"

#include <memory>
#include <psapi.h>
#include <tlhelp32.h>

#pragma comment(lib, "psapi.lib")

namespace ActiveBasic { namespace Common {

struct HandleDeleter
{
	typedef HANDLE pointer;

	void operator ()(pointer h) const
	{
		::CloseHandle(h);
	}
};

typedef std::unique_ptr<HANDLE, HandleDeleter> UniqueHandle;

boost::filesystem::path GetModuleFilePath(HANDLE hProcess, HMODULE hmod)
{	
	if (GetOSVersion().dwPlatformId == VER_PLATFORM_WIN32_NT)
	{
		WCHAR temp[MAX_PATH] = {};
		if (::GetModuleFileNameExW(hProcess, hmod, temp, MAX_PATH) != 0)
		{
			return temp;
		}
	}
	else
	{
		CHAR temp[MAX_PATH] = {};
		if (::GetModuleFileNameA(hmod, temp, MAX_PATH) != 0)
		{
			return boost::filesystem::path(Jenga::Common::ToWString(temp));
		}
	}
	throw std::runtime_error("GetModuleFileName or GetModuleFileNameEx failed");
}

namespace {

struct Toolhelp32Function
{
	decltype(::CreateToolhelp32Snapshot)* pCreateToolhelp32Snapshot;
	decltype(::Process32First)* pProcess32First;
	decltype(::Process32Next)* pProcess32Next;
	decltype(::Module32First)* pModule32First;
	decltype(::Module32Next)* pModule32Next;
};

Toolhelp32Function const& GetToolhelp32Function(HMODULE hmodKernel)
{
	static Toolhelp32Function const tf =
	{
		reinterpret_cast<decltype(::CreateToolhelp32Snapshot)*>(::GetProcAddress(hmodKernel, "CreateToolhelp32Snapshot")),
		reinterpret_cast<decltype(::Process32First)*>(::GetProcAddress(hmodKernel, "Process32First")),
		reinterpret_cast<decltype(::Process32Next)*>(::GetProcAddress(hmodKernel, "Process32Next")),
		reinterpret_cast<decltype(::Module32First)*>(::GetProcAddress(hmodKernel, "Module32First")),
		reinterpret_cast<decltype(::Module32Next)*>(::GetProcAddress(hmodKernel, "Module32Next")),
	};
	return tf;
}

UniqueHandle CreateToolhelp32SnapshotHelper(Toolhelp32Function const& tf, DWORD flags, DWORD processId = 0)
{
	UniqueHandle h(tf.pCreateToolhelp32Snapshot(flags, processId));
	if (h.get() == INVALID_HANDLE_VALUE)
	{
		h.release();
		return UniqueHandle();
	}
	return h;
}

} // unnamed namespace

std::vector<DWORD> GetProcesses()
{
	auto const& tf = GetToolhelp32Function(::GetModuleHandle(TEXT("kernel32")));
	if (tf.pCreateToolhelp32Snapshot)
	{
		UniqueHandle h = CreateToolhelp32SnapshotHelper(tf, TH32CS_SNAPPROCESS);
		if (h == nullptr)
		{
			return std::vector<DWORD>();
		}
		std::vector<DWORD> processes;
		PROCESSENTRY32 pe = {sizeof pe};
		if (tf.pProcess32First(h.get(), &pe))
		{
			do
			{
				processes.push_back(pe.th32ProcessID);
			} while (tf.pProcess32Next(h.get(), &pe));
		}
		return processes;
	}
	else
	{
		std::vector<DWORD> processes(8192);
		DWORD cbNeeded = 0;
		if (EnumProcesses(processes.data(), processes.size() * sizeof (DWORD), &cbNeeded))
		{
			processes.resize(cbNeeded / sizeof (DWORD));
			return processes;
		}
		return std::vector<DWORD>();
	}
}

boost::filesystem::path GetExecutableModulePath(DWORD processId)
{
	if (ActiveBasic::Common::Is9x())
	{
		auto const& tf = GetToolhelp32Function(::GetModuleHandle(TEXT("kernel32")));
		UniqueHandle h = CreateToolhelp32SnapshotHelper(tf, TH32CS_SNAPMODULE, processId);
		if (h == nullptr)
		{
			return boost::filesystem::path();
		}
		MODULEENTRY32 me = {sizeof me};
		if (tf.pModule32First(h.get(), &me))
		{
			return boost::filesystem::path(me.szExePath);
		}
		else
		{
			return boost::filesystem::path();
		}
	}
	else
	{
		UniqueHandle hProcess(::OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, processId));
		// PSAPIが使える環境はそちらを優先。
		WCHAR filename[MAX_PATH] = {};
		::GetModuleFileNameExW(hProcess.get(), nullptr, filename, MAX_PATH);
		return boost::filesystem::path(filename);
	}
}

std::vector<HMODULE> GetModules(DWORD processId)
{
	auto const& tf = GetToolhelp32Function(::GetModuleHandle(TEXT("kernel32")));
	if (tf.pCreateToolhelp32Snapshot)
	{
		UniqueHandle h = CreateToolhelp32SnapshotHelper(tf, TH32CS_SNAPMODULE, processId);
		if (h == nullptr)
		{
			return std::vector<HMODULE>();
		}
		std::vector<HMODULE> modules;
		MODULEENTRY32 me = {sizeof me};
		if (tf.pModule32First(h.get(), &me))
		{
			do
			{
				modules.push_back(me.hModule);
			} while (tf.pModule32Next(h.get(), &me));
		}
		return modules;
	}
	else
	{
		UniqueHandle hProcess(::OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, processId));
		const DWORD MAX_MODULE = 1024;
		std::vector<HMODULE> module(MAX_MODULE);
		DWORD cbReturned = 0;
		EnumProcessModules(hProcess.get(), module.data(), sizeof (HMODULE) * MAX_MODULE, &cbReturned);
		module.resize(cbReturned / sizeof (HMODULE));
		return module;
	}
}

}}