#include "animia/fd/win32.h"
#include "animia/util/win32.h"
#include "animia.h"

#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>

#include <fileapi.h>
#include <handleapi.h>
#include <libloaderapi.h>
#include <ntdef.h>
#include <psapi.h>
#include <shlobj.h>
#include <stringapiset.h>
#include <tlhelp32.h>
#include <windows.h>
#include <winternl.h>

/* This file is noticably more complex than Unix and Linux, and that's because
 * there is no "simple" way to get the paths of a file. In fact, this thing requires
 * you to use *internal functions* that can't even be linked to, hence why we have to
 * use GetProcAddress and such. What a mess.
 *
 * Speaking of which, because this file uses internal functions of the OS, it is not
 * guaranteed to work far into the future. However, it has worked since NT 6.0 (Vista)
 * at least, so it's unlikely to be changed much ever.
*/

/* SystemExtendedHandleInformation is only available in NT 5.1+ (XP and higher) and provides information for
 * 32-bit PIDs, unlike SystemHandleInformation
*/
constexpr SYSTEM_INFORMATION_CLASS SystemExtendedHandleInformation = static_cast<SYSTEM_INFORMATION_CLASS>(0x40);
constexpr NTSTATUS STATUS_INFO_LENGTH_MISMATCH = 0xC0000004UL;

/* this is filled in at runtime because it's not guaranteed to be (and isn't)
   constant between different versions of Windows */
static unsigned short file_type_index = 0;

struct SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX {
		PVOID Object;
		ULONG_PTR UniqueProcessId;
		HANDLE HandleValue;
		ACCESS_MASK GrantedAccess;
		USHORT CreatorBackTraceIndex;
		USHORT ObjectTypeIndex;
		ULONG HandleAttributes;
		ULONG Reserved;
};

struct SYSTEM_HANDLE_INFORMATION_EX {
		ULONG_PTR NumberOfHandles;
		ULONG_PTR Reserved;
		SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX Handles[1];
};

namespace animia::internal::win32 {

static HANDLE DuplicateHandle(HANDLE process_handle, HANDLE handle) {
	HANDLE dup_handle = nullptr;
	const bool result =
	    ::DuplicateHandle(process_handle, handle, ::GetCurrentProcess(), &dup_handle, 0, false, DUPLICATE_SAME_ACCESS);
	return result ? dup_handle : nullptr;
}

static PVOID GetNTDLLAddress(LPCSTR proc_name) {
	return reinterpret_cast<PVOID>(::GetProcAddress(::GetModuleHandleA("ntdll.dll"), proc_name));
}

static NTSTATUS QuerySystemInformation(SYSTEM_INFORMATION_CLASS cls, PVOID sysinfo, ULONG len, PULONG retlen) {
	static const auto func =
	    reinterpret_cast<decltype(::NtQuerySystemInformation)*>(GetNTDLLAddress("NtQuerySystemInformation"));
	return func(cls, sysinfo, len, retlen);
}

static NTSTATUS QueryObject(HANDLE handle, OBJECT_INFORMATION_CLASS cls, PVOID objinf, ULONG objinflen, PULONG retlen) {
	static const auto func = reinterpret_cast<decltype(::NtQueryObject)*>(GetNTDLLAddress("NtQueryObject"));
	return func(handle, cls, objinf, objinflen, retlen);
}

static std::vector<SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX> GetSystemHandleInformation() {
	std::vector<SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX> res;
	/* we should really put a cap on this */
	ULONG cb = 1 << 19;

	for (NTSTATUS status = STATUS_INFO_LENGTH_MISMATCH; status == STATUS_INFO_LENGTH_MISMATCH;) {
		/* why are we doing this? */
		status = STATUS_NO_MEMORY;

		SYSTEM_HANDLE_INFORMATION_EX* info = (SYSTEM_HANDLE_INFORMATION_EX*)malloc(cb *= 2);
		if (!info)
			continue;

		res.reserve(cb);

		status = QuerySystemInformation(SystemExtendedHandleInformation, info, cb, &cb);
		if (0 <= status) {
			ULONG_PTR handles = info->NumberOfHandles;
			if (handles) {
				res.reserve(res.size() + handles);

				SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX* entry = info->Handles;
				do {
					if (entry)
						res.push_back(*entry);
				} while (entry++, --handles);
			}
		}

		free(info);
	}

	return res;
}

static OBJECT_TYPE_INFORMATION QueryObjectTypeInfo(HANDLE handle) {
	OBJECT_TYPE_INFORMATION info;
	QueryObject(handle, ObjectTypeInformation, &info, sizeof(info), NULL);
	return info;
}

static std::string GetHandleType(HANDLE handle) {
	OBJECT_TYPE_INFORMATION info = QueryObjectTypeInfo(handle);
	return ToUtf8String(info.TypeName);
}

static std::string GetFinalPathNameByHandle(HANDLE handle) {
	std::wstring buffer;

	int result = ::GetFinalPathNameByHandleW(handle, NULL, 0, FILE_NAME_NORMALIZED | VOLUME_NAME_DOS);
	buffer.resize(result);
	::GetFinalPathNameByHandleW(handle, &buffer.front(), buffer.size(), FILE_NAME_NORMALIZED | VOLUME_NAME_DOS);
	buffer.resize(buffer.find('\0'));

	return ToUtf8String(buffer);
}

static bool IsFileHandle(HANDLE handle, unsigned short object_type_index) {
	if (file_type_index)
		return object_type_index == file_type_index;
	else if (!handle)
		return true;
	else if (GetHandleType(handle) == "File") {
		file_type_index = object_type_index;
		return true;
	}
	return false;
}

static bool IsFileMaskOk(ACCESS_MASK access_mask) {
	if (!(access_mask & FILE_READ_DATA))
		return false;

	if ((access_mask & FILE_APPEND_DATA) || (access_mask & FILE_WRITE_EA) || (access_mask & FILE_WRITE_ATTRIBUTES))
		return false;

	return true;
}

static bool IsFilePathOk(const std::string& path) {
	if (path.empty())
		return false;

	if (IsSystemDirectory(path))
		return false;

	const auto file_attributes = GetFileAttributesA(path.c_str());
	if ((file_attributes == INVALID_FILE_ATTRIBUTES) || (file_attributes & FILE_ATTRIBUTE_DIRECTORY))
		return false;

	return true;
}

bool EnumerateOpenProcesses(process_proc_t process_proc) {
	HANDLE hProcessSnap = ::CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hProcessSnap == INVALID_HANDLE_VALUE)
		return false;

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

	if (!::Process32First(hProcessSnap, &pe32))
		return false;

	if (!process_proc({pe32.th32ProcessID, pe32.szExeFile}))
		return false;

	while (::Process32Next(hProcessSnap, &pe32))
		if (!process_proc({pe32.th32ProcessID, pe32.szExeFile}))
			return false;

	::CloseHandle(hProcessSnap);

	return true;
}

/* this could be changed to being a callback, but... I'm too lazy right now :) */
bool EnumerateOpenFiles(const std::set<pid_t>& pids, open_file_proc_t open_file_proc) {
	if (!open_file_proc)
		return false;

	std::unordered_map<pid_t, Handle> proc_handles;

	for (const pid_t& pid : pids) {
		const HANDLE handle = ::OpenProcess(PROCESS_DUP_HANDLE, false, pid);
		if (handle != INVALID_HANDLE_VALUE)
			proc_handles[pid] = Handle(handle);
	}

	if (proc_handles.empty())
		return false;

	std::vector<SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX> info = GetSystemHandleInformation();

	for (const auto& h : info) {
		const pid_t pid = h.UniqueProcessId;
		if (!pids.count(pid))
			continue;

		if (!IsFileHandle(nullptr, h.ObjectTypeIndex))
			continue;

		if (!IsFileMaskOk(h.GrantedAccess))
			continue;

		Handle handle(DuplicateHandle(proc_handles[pid].get(), h.HandleValue));
		if (handle.get() == INVALID_HANDLE_VALUE)
			continue;

		if (GetFileType(handle.get()) != FILE_TYPE_DISK)
			continue;

		const std::string path = GetFinalPathNameByHandle(handle.get());
		if (!IsFilePathOk(path))
			continue;

		if (!open_file_proc({pid, path}))
			return false;
	}

	return true;
}

} // namespace animia::internal::win32