#pragma once
#include <windows.h>
#include <cstdint>

#include "../utils.hpp"
#include "../loadup.hpp"
#include "raw_driver.hpp"

#define MAP_PHYSICAL 0xC3502004
#define UNMAP_PHYSICAL 0xC3502008

#pragma pack (push, 1)
typedef struct _gdrv_t
{
	unsigned long	interface_type;
	unsigned long	bus;
	std::uintptr_t  phys_addr;
	unsigned long	io_space;
	unsigned long	size;
} gdrv_t, *pgdrv_t;
#pragma pack (pop)

namespace vdm
{
	inline HANDLE drv_handle;
	__forceinline auto load_drv() -> std::tuple<HANDLE, std::string, NTSTATUS>
	{
		const auto [result, key] =
			driver::load(
				vdm::raw_driver,
				sizeof(vdm::raw_driver)
			);

		if (result != STATUS_SUCCESS) 
			return { {}, {}, result };

		vdm::drv_handle = CreateFileA(
			"\\\\.\\GIO",
			GENERIC_READ | GENERIC_WRITE,
			NULL,
			NULL,
			OPEN_EXISTING,
			FILE_ATTRIBUTE_NORMAL,
			NULL
		);

		return { vdm::drv_handle, key, result };
	}

	__forceinline auto unload_drv(HANDLE drv_handle, std::string drv_key) -> NTSTATUS
	{
		if (!CloseHandle(drv_handle))
			return STATUS_FAIL_CHECK;

		return driver::unload(drv_key);
	}

	__forceinline bool read_phys(void* addr, void* buffer, std::size_t size)
	{
		gdrv_t in_buffer;
		in_buffer.bus = NULL;
		in_buffer.interface_type = NULL;
		in_buffer.phys_addr = reinterpret_cast<std::uintptr_t>(addr);
		in_buffer.io_space = NULL;
		in_buffer.size = size;

		void* out_buffer[2] = { 0 };
		unsigned long returned = 0;

		if (!DeviceIoControl(
			drv_handle,
			MAP_PHYSICAL,
			reinterpret_cast<void*>(&in_buffer),
			sizeof in_buffer,
			out_buffer,
			sizeof out_buffer,
			&returned, NULL
		))
			return false;

		__try
		{
			memcpy(buffer, out_buffer[0], size);
		}
		__except (EXCEPTION_EXECUTE_HANDLER)
		{}

		return DeviceIoControl(
			drv_handle,
			UNMAP_PHYSICAL,
			reinterpret_cast<void*>(&out_buffer[0]),
			sizeof out_buffer[0],
			out_buffer,
			sizeof out_buffer,
			&returned, NULL
		);
	}

	__forceinline bool write_phys(void* addr, void* buffer, std::size_t size)
	{
		gdrv_t in_buffer;
		in_buffer.bus = NULL;
		in_buffer.interface_type = NULL;
		in_buffer.phys_addr = reinterpret_cast<std::uintptr_t>(addr);
		in_buffer.io_space = NULL;
		in_buffer.size = size;

		void* out_buffer[2] = { 0 };
		unsigned long returned = 0;

		if (!DeviceIoControl(
			drv_handle,
			MAP_PHYSICAL,
			reinterpret_cast<void*>(&in_buffer),
			sizeof in_buffer,
			out_buffer,
			sizeof out_buffer,
			&returned, NULL
		))
			return false;

		__try
		{
			memcpy(out_buffer[0], buffer, size);
		}
		__except (EXCEPTION_EXECUTE_HANDLER)
		{}

		return DeviceIoControl(
			drv_handle,
			UNMAP_PHYSICAL,
			reinterpret_cast<void*>(&out_buffer[0]),
			sizeof out_buffer[0],
			out_buffer,
			sizeof out_buffer,
			&returned, NULL
		);
	}
}