#include "ntdll.h"

#include "common.h"
#include "context.h"
#include "errors.h"
#include "files.h"
#include "handles.h"
#include "heap.h"
#include "kernel32/internal.h"
#include "kernel32/processthreadsapi.h"
#include "modules.h"
#include "processes.h"
#include "strutil.h"
#include "types.h"

#include <cerrno>
#include <cstring>
#include <sys/stat.h>
#include <unistd.h>

#include <optional>

namespace {

struct PROCESS_BASIC_INFORMATION {
	NTSTATUS ExitStatus;
	PEB *PebBaseAddress;
	ULONG_PTR AffinityMask;
	LONG BasePriority;
	ULONG_PTR UniqueProcessId;
	ULONG_PTR InheritedFromUniqueProcessId;
};

struct ProcessHandleDetails {
	pid_t pid = -1;
	DWORD exitCode = STILL_ACTIVE;
	PEB *peb = nullptr;
	bool isCurrentProcess = false;
};

constexpr LONG kDefaultBasePriority = 8;

struct RTL_OSVERSIONINFOEXW : RTL_OSVERSIONINFOW {
	WORD wServicePackMajor;
	WORD wServicePackMinor;
	WORD wSuiteMask;
	BYTE wProductType;
	BYTE wReserved;
};

using PRTL_OSVERSIONINFOEXW = RTL_OSVERSIONINFOEXW *;

constexpr ULONG kOsMajorVersion = 6;
constexpr ULONG kOsMinorVersion = 2;
constexpr ULONG kOsBuildNumber = 0;
constexpr ULONG kOsPlatformId = 2;			// VER_PLATFORM_WIN32_NT
constexpr BYTE kProductTypeWorkstation = 1; // VER_NT_WORKSTATION

bool resolveProcessDetails(HANDLE processHandle, ProcessHandleDetails &details) {
	if (kernel32::isPseudoCurrentProcessHandle(processHandle)) {
		details.pid = getpid();
		details.exitCode = STILL_ACTIVE;
		details.peb = wibo::processPeb;
		details.isCurrentProcess = true;
		return true;
	}

	auto po = wibo::handles().getAs<ProcessObject>(processHandle);
	if (!po) {
		return false;
	}

	details.pid = po->pid;
	details.exitCode = po->exitCode;
	details.isCurrentProcess = po->pid == getpid();
	details.peb = details.isCurrentProcess ? wibo::processPeb : nullptr;
	return true;
}

std::string windowsImagePathFor(const ProcessHandleDetails &details) {
	if (details.isCurrentProcess && !wibo::guestExecutablePath.empty()) {
		return files::pathToWindows(files::canonicalPath(wibo::guestExecutablePath));
	}

	std::error_code ec;
	std::filesystem::path link = std::filesystem::path("/proc") / std::to_string(details.pid) / "exe";
	std::filesystem::path resolved = std::filesystem::read_symlink(link, ec);
	if (!ec) {
		return files::pathToWindows(files::canonicalPath(resolved));
	}
	return {};
}

} // namespace

namespace kernel32 {
BOOL WIN_FUNC SetEvent(HANDLE hEvent);
BOOL WIN_FUNC ResetEvent(HANDLE hEvent);
} // namespace kernel32

namespace ntdll {

constexpr LARGE_INTEGER FILE_WRITE_TO_END_OF_FILE = {.QuadPart = -1};
constexpr LARGE_INTEGER FILE_USE_FILE_POINTER_POSITION = {.QuadPart = -2};

PVOID CDECL memset(PVOID dest, int ch, SIZE_T count) {
	VERBOSE_LOG("ntdll::memset(%p, %i, %zu)\n", dest, ch, count);
	return std::memset(dest, ch, count);
}

NTSTATUS WINAPI NtReadFile(HANDLE FileHandle, HANDLE Event, PIO_APC_ROUTINE ApcRoutine, PVOID ApcContext,
						   PIO_STATUS_BLOCK IoStatusBlock, PVOID Buffer, ULONG Length, PLARGE_INTEGER ByteOffset,
						   PULONG Key) {
	HOST_CONTEXT_GUARD();
	DEBUG_LOG("NtReadFile(%p, %p, %p, %p, %p, %p, %u, %p, %p) ", FileHandle, Event, ApcRoutine, ApcContext,
			  IoStatusBlock, Buffer, Length, ByteOffset, Key);
	(void)ApcRoutine;
	(void)ApcContext;
	(void)Key;

	if (!IoStatusBlock) {
		return STATUS_INVALID_PARAMETER;
	}
	IoStatusBlock->Information = 0;

	auto file = wibo::handles().getAs<FileObject>(FileHandle);
	if (!file || !file->valid()) {
		IoStatusBlock->Status = STATUS_INVALID_HANDLE;
		IoStatusBlock->Information = 0;
		return STATUS_INVALID_HANDLE;
	}

	bool useOverlapped = file->overlapped;
	bool useCurrentFilePosition = (ByteOffset == nullptr);
	if (!useCurrentFilePosition && ByteOffset->QuadPart == FILE_USE_FILE_POINTER_POSITION.QuadPart) {
		useCurrentFilePosition = true;
	}

	std::optional<off_t> offset;
	if (!useCurrentFilePosition) {
		offset = static_cast<off_t>(ByteOffset->QuadPart);
	}

	if (useOverlapped && useCurrentFilePosition) {
		IoStatusBlock->Status = STATUS_INVALID_PARAMETER;
		IoStatusBlock->Information = 0;
		return STATUS_INVALID_PARAMETER;
	}

	Pin<kernel32::EventObject> ev;
	if (Event) {
		ev = wibo::handles().getAs<kernel32::EventObject>(Event);
		if (!ev) {
			IoStatusBlock->Status = STATUS_INVALID_HANDLE;
			IoStatusBlock->Information = 0;
			return STATUS_INVALID_HANDLE;
		}
		ev->reset();
	}

	bool updateFilePointer = !useOverlapped;
	auto io = files::read(file.get(), Buffer, Length, offset, updateFilePointer);
	NTSTATUS status = STATUS_SUCCESS;
	if (io.unixError != 0) {
		status = wibo::statusFromErrno(io.unixError);
	} else if (io.reachedEnd && io.bytesTransferred == 0) {
		status = file->isPipe ? STATUS_PIPE_BROKEN : STATUS_END_OF_FILE;
	}

	IoStatusBlock->Status = status;
	IoStatusBlock->Information = static_cast<ULONG_PTR>(io.bytesTransferred);

	if (ev && (status == STATUS_SUCCESS || status == STATUS_END_OF_FILE)) {
		ev->set();
	}

	DEBUG_LOG("-> 0x%x\n", status);
	return status;
}

NTSTATUS WINAPI NtWriteFile(HANDLE FileHandle, HANDLE Event, PIO_APC_ROUTINE ApcRoutine, PVOID ApcContext,
							PIO_STATUS_BLOCK IoStatusBlock, PVOID Buffer, ULONG Length, PLARGE_INTEGER ByteOffset,
							PULONG Key) {
	HOST_CONTEXT_GUARD();
	DEBUG_LOG("NtWriteFile(%p, %p, %p, %p, %p, %p, %u, %p, %p) ", FileHandle, Event, ApcRoutine, ApcContext,
			  IoStatusBlock, Buffer, Length, ByteOffset, Key);
	(void)ApcRoutine;
	(void)ApcContext;
	(void)Key;

	if (!IoStatusBlock) {
		return STATUS_INVALID_PARAMETER;
	}
	IoStatusBlock->Information = 0;

	auto file = wibo::handles().getAs<FileObject>(FileHandle);
	if (!file || !file->valid()) {
		IoStatusBlock->Status = STATUS_INVALID_HANDLE;
		return STATUS_INVALID_HANDLE;
	}

	bool useOverlapped = file->overlapped;
	bool useCurrentFilePosition = (ByteOffset == nullptr);
	bool writeToEndOfFile = false;
	if (ByteOffset) {
		if (ByteOffset->QuadPart == FILE_USE_FILE_POINTER_POSITION.QuadPart) {
			useCurrentFilePosition = true;
		} else if (ByteOffset->QuadPart == FILE_WRITE_TO_END_OF_FILE.QuadPart) {
			writeToEndOfFile = true;
		}
	}

	std::optional<off_t> offset;
	if (!useCurrentFilePosition && !writeToEndOfFile) {
		offset = static_cast<off_t>(ByteOffset->QuadPart);
	}

	if (useOverlapped && useCurrentFilePosition) {
		IoStatusBlock->Status = STATUS_INVALID_PARAMETER;
		return STATUS_INVALID_PARAMETER;
	}

	Pin<kernel32::EventObject> ev;
	if (Event) {
		ev = wibo::handles().getAs<kernel32::EventObject>(Event);
		if (!ev) {
			IoStatusBlock->Status = STATUS_INVALID_HANDLE;
			return STATUS_INVALID_HANDLE;
		}
		ev->reset();
	}

	bool updateFilePointer = file->isPipe ? true : !useOverlapped;

	if (writeToEndOfFile && !offset.has_value()) {
		if (!file->isPipe) {
			struct stat st{};
			if (fstat(file->fd, &st) != 0) {
				int err = errno ? errno : EIO;
				NTSTATUS status = wibo::statusFromErrno(err);
				IoStatusBlock->Status = status;
				return status;
			}
			offset = static_cast<off_t>(st.st_size);
		}
	}

	auto io = files::write(file.get(), Buffer, static_cast<size_t>(Length), offset, updateFilePointer);
	NTSTATUS status = STATUS_SUCCESS;
	if (io.unixError != 0) {
		status = wibo::statusFromErrno(io.unixError);
	}

	IoStatusBlock->Status = status;
	IoStatusBlock->Information = static_cast<ULONG_PTR>(io.bytesTransferred);

	if (ev && status == STATUS_SUCCESS) {
		ev->set();
	}

	DEBUG_LOG("-> 0x%x\n", status);
	return status;
}

NTSTATUS WINAPI NtAllocateVirtualMemory(HANDLE ProcessHandle, PVOID *BaseAddress, ULONG_PTR ZeroBits,
										PSIZE_T RegionSize, ULONG AllocationType, ULONG Protect) {
	HOST_CONTEXT_GUARD();
	DEBUG_LOG("NtAllocateVirtualMemory(%p, %p, %lu, %p, %lu, %lu) ", ProcessHandle, BaseAddress, ZeroBits, RegionSize,
			  AllocationType, Protect);
	if (ProcessHandle != (HANDLE)-1) {
		DEBUG_LOG("-> 0x%x\n", STATUS_INVALID_HANDLE);
		return STATUS_INVALID_HANDLE;
	}
	if (ZeroBits != 0 || BaseAddress == nullptr || RegionSize == nullptr) {
		DEBUG_LOG("-> 0x%x\n", STATUS_INVALID_PARAMETER);
		return STATUS_INVALID_PARAMETER;
	}

	wibo::heap::VmStatus vmStatus =
		wibo::heap::virtualAlloc(reinterpret_cast<void **>(BaseAddress), reinterpret_cast<std::size_t *>(RegionSize),
								 static_cast<DWORD>(AllocationType), static_cast<DWORD>(Protect));
	if (vmStatus != wibo::heap::VmStatus::Success) {
		NTSTATUS status = wibo::heap::ntStatusFromVmStatus(vmStatus);
		DEBUG_LOG("-> 0x%x\n", status);
		return status;
	}

	DEBUG_LOG("-> 0x%x\n", STATUS_SUCCESS);
	return STATUS_SUCCESS;
}

NTSTATUS WINAPI NtProtectVirtualMemory(HANDLE ProcessHandle, PVOID *BaseAddress, PSIZE_T NumberOfBytesToProtect,
									   ULONG NewAccessProtection, PULONG OldAccessProtection) {
	HOST_CONTEXT_GUARD();
	DEBUG_LOG("NtProtectVirtualMemory(%p, %p, %p, %lu, %p) ", ProcessHandle, BaseAddress, NumberOfBytesToProtect,
			  NewAccessProtection, OldAccessProtection);
	if (ProcessHandle != (HANDLE)-1) {
		DEBUG_LOG("-> 0x%x\n", STATUS_INVALID_HANDLE);
		return STATUS_INVALID_HANDLE;
	}
	if (BaseAddress == nullptr || NumberOfBytesToProtect == nullptr) {
		DEBUG_LOG("-> 0x%x\n", STATUS_INVALID_PARAMETER);
		return STATUS_INVALID_PARAMETER;
	}

	void *base = *BaseAddress;
	std::size_t length = static_cast<std::size_t>(*NumberOfBytesToProtect);
	wibo::heap::VmStatus vmStatus =
		wibo::heap::virtualProtect(base, length, static_cast<DWORD>(NewAccessProtection), OldAccessProtection);
	if (vmStatus != wibo::heap::VmStatus::Success) {
		NTSTATUS status = wibo::heap::ntStatusFromVmStatus(vmStatus);
		DEBUG_LOG("-> 0x%x\n", status);
		return status;
	}

	DEBUG_LOG("-> 0x%x\n", STATUS_SUCCESS);
	return STATUS_SUCCESS;
}

NTSTATUS WINAPI RtlGetVersion(PRTL_OSVERSIONINFOW lpVersionInformation) {
	HOST_CONTEXT_GUARD();
	DEBUG_LOG("RtlGetVersion(%p) ", lpVersionInformation);
	if (!lpVersionInformation) {
		DEBUG_LOG("-> 0x%x\n", STATUS_INVALID_PARAMETER);
		return STATUS_INVALID_PARAMETER;
	}

	ULONG size = lpVersionInformation->dwOSVersionInfoSize;
	if (size < sizeof(RTL_OSVERSIONINFOW)) {
		DEBUG_LOG("-> 0x%x\n", STATUS_INVALID_PARAMETER);
		return STATUS_INVALID_PARAMETER;
	}

	std::memset(lpVersionInformation, 0, static_cast<size_t>(size));
	lpVersionInformation->dwOSVersionInfoSize = size;
	lpVersionInformation->dwMajorVersion = kOsMajorVersion;
	lpVersionInformation->dwMinorVersion = kOsMinorVersion;
	lpVersionInformation->dwBuildNumber = kOsBuildNumber;
	lpVersionInformation->dwPlatformId = kOsPlatformId;

	if (size >= sizeof(RTL_OSVERSIONINFOEXW)) {
		auto extended = reinterpret_cast<PRTL_OSVERSIONINFOEXW>(lpVersionInformation);
		extended->wProductType = kProductTypeWorkstation;
	}

	DEBUG_LOG("-> 0x%x\n", STATUS_SUCCESS);
	return STATUS_SUCCESS;
}

NTSTATUS WINAPI NtQueryInformationProcess(HANDLE ProcessHandle, PROCESSINFOCLASS ProcessInformationClass,
										  PVOID ProcessInformation, ULONG ProcessInformationLength,
										  PULONG ReturnLength) {
	HOST_CONTEXT_GUARD();
	DEBUG_LOG("NtQueryInformationProcess(%p, %u, %p, %u, %p) ", ProcessHandle, ProcessInformationClass,
			  ProcessInformation, ProcessInformationLength, ReturnLength);
	if (!ProcessInformation) {
		DEBUG_LOG("-> 0x%x\n", STATUS_INVALID_PARAMETER);
		return STATUS_INVALID_PARAMETER;
	}

	ProcessHandleDetails details{};
	if (!resolveProcessDetails(ProcessHandle, details)) {
		DEBUG_LOG("-> 0x%x\n", STATUS_INVALID_HANDLE);
		return STATUS_INVALID_HANDLE;
	}

	switch (ProcessInformationClass) {
	case ProcessBasicInformation: {
		size_t required = sizeof(PROCESS_BASIC_INFORMATION);
		if (ReturnLength) {
			*ReturnLength = static_cast<ULONG>(required);
		}
		if (ProcessInformationLength < required) {
			DEBUG_LOG("-> 0x%x\n", STATUS_INFO_LENGTH_MISMATCH);
			return STATUS_INFO_LENGTH_MISMATCH;
		}

		auto *info = reinterpret_cast<PROCESS_BASIC_INFORMATION *>(ProcessInformation);
		std::memset(info, 0, sizeof(*info));
		info->ExitStatus = static_cast<NTSTATUS>(details.exitCode);
		info->PebBaseAddress = details.peb;
		DWORD_PTR processMask = 0;
		DWORD_PTR systemMask = 0;
		if (kernel32::GetProcessAffinityMask(ProcessHandle, &processMask, &systemMask)) {
			info->AffinityMask = static_cast<ULONG_PTR>(processMask == 0 ? 1 : processMask);
		} else {
			info->AffinityMask = 1;
		}
		info->BasePriority = kDefaultBasePriority;
		info->UniqueProcessId = static_cast<ULONG_PTR>(details.pid);
		if (details.isCurrentProcess) {
			info->InheritedFromUniqueProcessId = static_cast<ULONG_PTR>(getppid());
		} else {
			info->InheritedFromUniqueProcessId = static_cast<ULONG_PTR>(getpid());
		}
		DEBUG_LOG("-> 0x%x\n", STATUS_SUCCESS);
		return STATUS_SUCCESS;
	}
	case ProcessWow64Information: {
		size_t required = sizeof(ULONG_PTR);
		if (ReturnLength) {
			*ReturnLength = static_cast<ULONG>(required);
		}
		if (ProcessInformationLength < required) {
			DEBUG_LOG("-> 0x%x\n", STATUS_INFO_LENGTH_MISMATCH);
			return STATUS_INFO_LENGTH_MISMATCH;
		}
		auto *value = reinterpret_cast<ULONG_PTR *>(ProcessInformation);
		*value = 0;
		DEBUG_LOG("-> 0x%x\n", STATUS_SUCCESS);
		return STATUS_SUCCESS;
	}
	case ProcessImageFileName: {
		size_t minimum = sizeof(UNICODE_STRING);
		if (ProcessInformationLength < minimum) {
			if (ReturnLength) {
				*ReturnLength = static_cast<ULONG>(minimum);
			}
			DEBUG_LOG("-> 0x%x\n", STATUS_INFO_LENGTH_MISMATCH);
			return STATUS_INFO_LENGTH_MISMATCH;
		}

		std::string imagePath = windowsImagePathFor(details);
		DEBUG_LOG("  NtQueryInformationProcess image path: %s\n", imagePath.c_str());
		auto widePath = stringToWideString(imagePath.c_str());
		size_t stringBytes = widePath.size() * sizeof(uint16_t);
		size_t required = sizeof(UNICODE_STRING) + stringBytes;
		if (ReturnLength) {
			*ReturnLength = static_cast<ULONG>(required);
		}
		if (ProcessInformationLength < required) {
			DEBUG_LOG("-> 0x%x\n", STATUS_INFO_LENGTH_MISMATCH);
			return STATUS_INFO_LENGTH_MISMATCH;
		}

		auto *unicode = reinterpret_cast<UNICODE_STRING *>(ProcessInformation);
		auto *buffer =
			reinterpret_cast<uint16_t *>(reinterpret_cast<uint8_t *>(ProcessInformation) + sizeof(UNICODE_STRING));
		std::memcpy(buffer, widePath.data(), stringBytes);
		size_t characterCount = widePath.empty() ? 0 : widePath.size() - 1;
		unicode->Length = static_cast<unsigned short>(characterCount * sizeof(uint16_t));
		unicode->MaximumLength = static_cast<unsigned short>(widePath.size() * sizeof(uint16_t));
		unicode->Buffer = buffer;
		DEBUG_LOG("-> 0x%x\n", STATUS_SUCCESS);
		return STATUS_SUCCESS;
	}
	default:
		DEBUG_LOG("-> 0x%x\n", STATUS_INVALID_INFO_CLASS);
		return STATUS_INVALID_INFO_CLASS;
	}
}

} // namespace ntdll

#include "ntdll_trampolines.h"

extern const wibo::ModuleStub lib_ntdll = {
	(const char *[]){
		"ntdll",
		nullptr,
	},
	ntdllThunkByName,
	nullptr,
};