#ifndef EFI_CPP
#define EFI_CPP

#include <core/inc.hpp>

#define offsetof(s,m) ((::size_t)&reinterpret_cast<char const volatile&>((((s*)0)->m)))

auto efi_hook::efi_t::is_dtb_invalid( std::uint64_t cr3 ) -> bool {
	return ( cr3 >> 0x38 ) == 0x40;
}

auto efi_hook::efi_t::create_payload_handle( ) -> bool {
	UNICODE_STRING deviceName;
	OBJECT_ATTRIBUTES objectAttributes;
	IO_STATUS_BLOCK ioStatusBlock;
	NTSTATUS status;

	RtlInitUnicodeString( &deviceName, hash_str( L"\\Device\\KsecDD" ) );

	InitializeObjectAttributes( &objectAttributes, &deviceName, 0, NULL, NULL );

	status = NtOpenFile(
		&this->payload_handle,
		GENERIC_READ | GENERIC_WRITE,
		&objectAttributes,
		&ioStatusBlock,
		FILE_SHARE_READ | FILE_SHARE_WRITE,
		0
	);

	if ( this->payload_handle != INVALID_HANDLE_VALUE ) {
		return true;
	}

	return false;
}

auto efi_hook::efi_t::is_payload_available( ) -> bool {
	auto data = is_mapped_invoke( );

	data.is_mapped = 0;

	this->send_payload_cmd(
		&data,
		invoke_is_mapped
	);

	return data.is_mapped;
}


auto efi_hook::efi_t::setup( const std::wstring& process_name ) -> setup_error_code {

	m_log->set_title( hash_str( "calamari-fortnite" ) );
	m_log->set_text( hash_str( "initializing calamari auto-updating bootstrapper..." ) );

	bool enabled = false;
	auto status = RtlAdjustPrivilege( 22L, true, false, reinterpret_cast< BOOLEAN* >( &enabled ) );

	if ( !nt_success( status ) )
		return setup_error_code::failed_to_gain_privilege;

	m_log->set_lab_execution(
		hash_str( "PDB Installation Execution" ),
		hash_str( "efi_hook::efi_t::efi_t()" )
	);

	m_log->set_seperator( );
	m_log->set_text( hash_str( "finding kernel imports..." ) );

	this->mm_copy_memory = this->get_kernel_import( hash_str( "ntoskrnl.exe" ), hash_str( "MmCopyMemory" ) );
	this->mm_map_space = this->get_kernel_import( hash_str( "ntoskrnl.exe" ), hash_str( "MmMapIoSpace" ) );
	this->mm_unmap_space = this->get_kernel_import( hash_str( "ntoskrnl.exe" ), hash_str( "MmUnmapIoSpace" ) );
	this->mm_physical_ranges = this->get_kernel_import( hash_str( "ntoskrnl.exe" ), hash_str( "MmGetPhysicalMemoryRanges" ) );
	this->ps_initial_system_process = this->get_kernel_import( hash_str( "ntoskrnl.exe" ), hash_str( "PsInitialSystemProcess" ) );
	this->ps_get_process_id = this->get_kernel_import( hash_str( "ntoskrnl.exe" ), hash_str( "PsGetCurrentProcessId" ) );
	this->ps_get_current_process = this->get_kernel_import( hash_str( "ntoskrnl.exe" ), hash_str( "PsGetCurrentProcess" ) );
	this->memcpy = this->get_kernel_import( hash_str( "ntoskrnl.exe" ), hash_str( "memcpy" ) );
	this->ex_allocate_pool = this->get_kernel_import( hash_str( "ntoskrnl.exe" ), hash_str( "ExAllocatePool" ) );
	this->ex_free_pool = this->get_kernel_import( hash_str( "ntoskrnl.exe" ), hash_str( "ExFreePool" ) );

	m_log->set_text( hash_str( "successfully found all imports." ) );
	m_log->set_seperator( );

	m_log->set_text( hash_str( "installing PDBs..." ) );

	srand( time( NULL ) );

	const auto& ntos_pdb_path = g_pdb->download_pdb_to_temp( hash_str( "\\System32\\ntoskrnl.exe" ) );

	if ( ntos_pdb_path == hash_str( "NONE" ) )
		return setup_error_code::pdb_failure;

	m_log->set_text_success( hash_str( "ntoskrnl.exe PDB downloaded at: %s" ), ntos_pdb_path.c_str( ) );

	auto ntos_pdb_object = pdb_interface::pdb( );

	const auto& ntos_pdb_status = g_pdb->load_pdb( ntos_pdb_path, &ntos_pdb_object );

	if ( !ntos_pdb_status )
		return setup_error_code::pdb_failure;

	process_active_process_links_offset = g_pdb->get_struct_property_offset(
		&ntos_pdb_object, hash_str( "_EPROCESS" ),
		hash_str( L"ActiveProcessLinks" )
	);

	process_unique_id_offset = g_pdb->get_struct_property_offset(
		&ntos_pdb_object, hash_str( "_EPROCESS" ),
		hash_str( L"UniqueProcessId" )
	);

	process_section_base_address_offset = g_pdb->get_struct_property_offset(
		&ntos_pdb_object, hash_str( "_EPROCESS" ),
		hash_str( L"SectionBaseAddress" )
	);

	thread_list_entry_offset = g_pdb->get_struct_property_offset(
		&ntos_pdb_object, hash_str( "_ETHREAD" ),
		hash_str( L"ThreadListEntry" )
	);

	thread_win32startaddress_offset = g_pdb->get_struct_property_offset(
		&ntos_pdb_object, hash_str( "_ETHREAD" ),
		hash_str( L"Win32StartAddress" )
	);

	thread_startaddress_offset = g_pdb->get_struct_property_offset(
		&ntos_pdb_object, hash_str( "_ETHREAD" ),
		hash_str( L"StartAddress" )
	);

	mm_allocate_independent_pages_offset = g_pdb->get_rva( &ntos_pdb_object, hash_str( "MmAllocateIndependentPages" ) );
	mm_free_independent_pages_offset = g_pdb->get_rva( &ntos_pdb_object, hash_str( "MmFreeIndependentPages" ) );
	mm_set_page_protection_offset = g_pdb->get_rva( &ntos_pdb_object, hash_str( "MmSetPageProtection" ) );

	m_log->set_text_point( hash_str( "thread_list_entry_offset: 0x%llx" ), thread_list_entry_offset );
	m_log->set_text_point( hash_str( "thread_win32startaddress_offset: 0x%llx" ), thread_win32startaddress_offset );
	m_log->set_text_point( hash_str( "thread_startaddress_offset: 0x%llx" ), thread_startaddress_offset );
	m_log->set_text_point( hash_str( "process_section_base_address_offset: 0x%llx" ), process_section_base_address_offset );
	m_log->set_text_point( hash_str( "process_pid_offset: 0x%llx" ), process_unique_id_offset );
	m_log->set_text_point( hash_str( "process_active_process_links_offset: 0x%llx" ), process_active_process_links_offset );
	m_log->set_text_point( hash_str( "mm_allocate_independent_pages_offset: 0x%llx" ), mm_allocate_independent_pages_offset );
	m_log->set_text_point( hash_str( "mm_free_independent_pages_offset: 0x%llx" ), mm_free_independent_pages_offset );
	m_log->set_text_point( hash_str( "mm_set_page_protection_offset: 0x%llx" ), mm_set_page_protection_offset );

	DeleteFileA( ntos_pdb_path.c_str( ) );
	g_pdb->unload_pdb( &ntos_pdb_object );

	//const auto& win32k_pdb_path = g_pdb->download_pdb_to_temp( hash_str( "\\System32\\win32kfull.sys" ) );

	//if ( win32k_pdb_path == hash_str( "NONE" ) )
	//	return setup_error_code::pdb_failure;

	//m_log->set_text_success( hash_str( "win32kfull.sys PDB downloaded at: %s" ), win32k_pdb_path.c_str( ) );

	//auto win32k_pdb_object = pdb_interface::pdb( );

	//const auto& win32k_pdb_status = g_pdb->load_pdb( win32k_pdb_path, &win32k_pdb_object );

	//if ( !win32k_pdb_status )
	//	return setup_error_code::pdb_failure;

	//mm_get_prop_offset = g_pdb->get_rva( &win32k_pdb_object, hash_str( "NtUserDispatchMessage" ) ) - 0x70;

	//m_log->set_text_point( hash_str( "mm_get_prop_offset: 0x%llx" ), mm_get_prop_offset );

	//g_pdb->unload_pdb( &win32k_pdb_object );
	//DeleteFileA( win32k_pdb_path.c_str( ) );


	m_log->set_lab_execution(
		hash_str( "EFI Execution" ),
		hash_str( "efi_hook::efi_t::setup()" )
	);

	m_log->set_seperator( );

	const auto& game_pid = get_process_pid( process_name.c_str( ) );

	if ( !game_pid )
		return setup_error_code::efi_unavailable;

	m_log->set_text_point(
		hash_str( "game pid: %i" ),
		game_pid
	);

	process_id = std::move( game_pid );
	ntos_base = get_kernel_image( hash_str( "ntoskrnl.exe" ) );

	if ( !ntos_base )
		return setup_error_code::efi_unavailable;

	const auto& payload_handle_status = this->create_payload_handle( );

	if ( !payload_handle_status ) {
		m_log->set_text_point( hash_str( "failed to create a handle with the payload target object.\n" ) );
		return setup_error_code::payload_communication_failure;
	}

	const auto& payload_availability_status = this->is_payload_available( );

	if ( !payload_availability_status ) {
		m_log->set_text_point( hash_str( "payload is not available.\n" ) );
		return setup_error_code::payload_mapping_failure;
	}

	//g_vm->hide_process( );

	const auto& payload_dtb_status = this->patch_directory_table_base( );

	if ( !payload_dtb_status ) {
		m_log->set_text_point( hash_str( "failed to patch dtb.\n" ) );
		return setup_error_code::invalid_directory_table_base;
	}

	cr3_header = payload_dtb_status;

	m_log->set_text_point(
		hash_str( "cr3: 0x%llx" ),
		cr3_header
	);

	const auto& base_address = get_image_base( );
	if ( !base_address )
		return setup_error_code::invalid_image_base;

	m_log->set_text_point(
		hash_str( "base address: 0x%llx" ),
		base_address
	);

	this->image_base = std::move( base_address );

	return setup_error_code::system_operational;
}
//
//auto efi_hook::efi_t::attach( ) -> attach_error_code
//{
//	m_log->set_lab_execution(
//		hash_str( "EFI Execution" ),
//		hash_str( "efi_hook::efi_t::setup()" )
//	);
//
//	m_log->set_seperator( );
//
//	const auto& game_pid = get_process_pid( process_name.c_str( ) );
//
//	if ( !game_pid )
//		return attach_error_code::efi_unavailable;
//
//	m_log->set_text_point(
//		hash_str( "game pid: %i" ),
//		game_pid
//	);
//
//	process_id = std::move( game_pid );
//	ntos_base = get_kernel_image( hash_str( "ntoskrnl.exe" ) );
//
//	if ( !ntos_base )
//		return attach_error_code::efi_unavailable;
//
//	const auto& payload_dtb_status = this->patch_directory_table_base( );
//
//	if ( !payload_dtb_status ) {
//		m_log->set_text_point( hash_str( "failed to patch dtb.\n" ) );
//		return attach_error_code::invalid_directory_table_base;
//	}
//
//	cr3_header = payload_dtb_status;
//
//	m_log->set_text_point(
//		hash_str( "cr3: 0x%llx" ),
//		cr3_header
//	);
//
//	const auto& base_address = get_image_base( );
//	if ( !base_address )
//		return attach_error_code::invalid_image_base;
//
//	m_log->set_text_point(
//		hash_str( "base address: 0x%llx" ),
//		base_address
//	);
//
//	this->image_base = std::move( base_address );
//
//	return attach_error_code::system_operational;
//}

auto efi_hook::efi_t::find_game( const std::wstring& process_name ) -> game_error_code
{
	if ( process_name.empty( ) ) return game_error_code::empty_process_name;

	while ( !get_process_pid( process_name ) ) {
		Sleep( 100 );
	}

	MessageBox( NULL, hash_str( L"It looks like everything is set up!\n\nPlease start the game and press OK inside the lobby.\n\nIf you encounter any issues, visit:\nhttps://setup.calamari.lol/troubleshooting" ), hash_str( L"Calamari" ), MB_OK | MB_ICONINFORMATION );

	return game_error_code::system_operational;
}


auto efi_hook::efi_t::send_payload_cmd(
	void* data,
	efi_hook::requests parameters ) -> bool {

	auto request = invoke_data( );

	request.unique = invoke_unique;
	request.data = data;
	request.code = parameters;

	const auto status = DeviceIoControl(
		this->payload_handle,
		PAYLOAD_CONTROL_CODE,
		&request,
		sizeof( request ),
		&request,
		sizeof( request ),
		nullptr,
		nullptr
	);

	return true;
}

auto efi_hook::efi_t::read_physical(
	const std::uint64_t address,
	void* buffer,
	const std::uint64_t size,
	bool large_page ) -> bool {

	if ( !memory::is_valid( address ) )
		return false;

	auto data = efi_hook::read_invoke( );

	data.pid = this->process_id;
	data.address = address;
	data.buffer = buffer;
	data.size = size;

	this->send_payload_cmd(
		&data,
		invoke_read
	);

	return true;
}

auto efi_hook::efi_t::get_image_base( ) -> std::uint64_t
{
	auto data = efi_hook::base_invoke( );

	data.pid = this->process_id;

	this->send_payload_cmd(
		&data,
		invoke_base
	);

	return data.handle;
}

auto efi_hook::efi_t::get_module_base( const wchar_t* module_name ) -> std::uint64_t
{
	auto data = efi_hook::peb_invoke( );

	data.pid = this->process_id;

	this->send_payload_cmd(
		&data,
		invoke_peb
	);

	auto peb = this->read<PEB64>( data.handle );
	if ( !peb.Ldr )
		return 0;

	auto ldr = this->read<MOD_LDR_DATA>( ( std::uint64_t ) ( peb.Ldr ) );
	if ( !ldr.InLoadOrderLinks.Flink )
		return 0;

	auto mod_entry = PMOD_DATA_TABLE_ENTRY( );

	for ( auto current = ldr.InLoadOrderLinks.Flink; ( PUCHAR ) current != ( PUCHAR ) peb.Ldr + offsetof( MOD_LDR_DATA, MOD_LDR_DATA::InLoadOrderLinks ); current = ( PLIST_ENTRY ) mod_entry->InLoadOrderLinks.Flink )
	{
		auto current_data = read<MOD_DATA_TABLE_ENTRY>( ( ULONG64 ) current );
		mod_entry = CONTAINING_RECORD( &current_data, MOD_DATA_TABLE_ENTRY, InLoadOrderLinks );

		if ( mod_entry->FullDllName.Buffer )
		{
			wchar_t name_buffer [ 256 ] = { 0 };
			read_physical( ( std::uint64_t ) mod_entry->FullDllName.Buffer, &name_buffer, sizeof( name_buffer ) );

			if ( wcsstr( name_buffer, module_name ) )
				return std::uint64_t( mod_entry->DllBase );
		}
	}

	return 0;
}

auto efi_hook::efi_t::read_large_array( std::uint64_t address, void* buffer, std::size_t size ) -> bool {

	if ( !memory::is_valid( address ) || !process_id )
		return false;

	constexpr std::uint64_t chunk_size = 512;
	std::uint64_t remaining_size = size;
	std::uint64_t current_address = address;
	auto current_buffer = reinterpret_cast< std::uint64_t >( buffer );

	while ( remaining_size > 0 )
	{
		std::uint64_t page_boundary = ( current_address & ~0xFFF ) + 0x1000;
		std::uint64_t bytes_to_read = std::min( { chunk_size, remaining_size, page_boundary - current_address } );
		unsigned char temp_buffer [ chunk_size ];

		if ( !read_physical( current_address, temp_buffer, bytes_to_read ) )
		{
			return false;
		}

		std::memcpy( reinterpret_cast< void* >( current_buffer ), temp_buffer, bytes_to_read );

		remaining_size -= bytes_to_read;
		current_address += bytes_to_read;
		current_buffer += bytes_to_read;
	}

	return true;

}

auto efi_hook::efi_t::write_physical(
	const std::uint64_t address,
	void* buffer,
	const std::uint64_t size ) -> bool {

	auto data = efi_hook::write_invoke( );

	data.pid = this->process_id;
	data.address = address;
	data.buffer = buffer;
	data.size = size;

	this->send_payload_cmd(
		&data,
		invoke_write
	);

	return true;
}

auto efi_hook::efi_t::patch_directory_table_base( ) -> std::uint64_t {

	auto data = efi_hook::dtb_invoke( );

	data.pid = this->process_id;
	data.handle = 0;

	this->send_payload_cmd(
		&data,
		invoke_resolve_dtb
	);

	return data.handle;
}

auto efi_hook::efi_t::get_peb( ) -> std::uint64_t {

	auto data = efi_hook::peb_invoke( );

	data.pid = this->process_id;
	data.handle = 0;

	this->send_payload_cmd(
		&data,
		invoke_peb
	);

	return data.handle;
}

auto efi_hook::efi_t::hide_process( ) -> bool {

	auto data = efi_hook::hide_process_invoke( );

	data.name = this->get_current_executable_name( );
	data.status = 0;

	this->send_payload_cmd(
		&data,
		invoke_hide_process
	);

	return data.status;
}

auto efi_hook::efi_t::hide_window_attributes( std::uint64_t window_handle ) -> bool {

	auto data = efi_hook::hide_attributes_invoke( );

	data.handle = window_handle;
	data.status = 0;

	this->send_payload_cmd(
		&data,
		invoke_hide_window
	);

	return data.status;
}

auto efi_hook::efi_t::get_kernel_import( const std::string& image, const std::string& module_name ) -> const std::uint64_t {
	const auto ntoskrnl_base = this->get_kernel_image( image.c_str( ) );
	if ( !ntoskrnl_base ) {
		m_log->set_text_point( hash_str( "failed to get ntoskrnl.exe" ) );
		return 0;
	}

	const auto func_addr = this->get_kernel_export(
		ntoskrnl_base,
		module_name.c_str( )
	);

	if ( !func_addr ) {
		m_log->set_text_point( hash_str( "failed to resolve import from ntoskrnl.exe" ) );
		return 0;
	}

	m_log->set_text_point( hash_str( "%s: 0x%llx" ), module_name.c_str( ), func_addr );

	return func_addr;
}

auto efi_hook::efi_t::get_process_pid( const std::wstring& proc_name ) -> const std::uint32_t {
	PROCESSENTRY32 proc_info;
	proc_info.dwSize = sizeof( proc_info );

	HANDLE proc_snapshot = CreateToolhelp32Snapshot( TH32CS_SNAPPROCESS, NULL );
	if ( proc_snapshot == INVALID_HANDLE_VALUE ) {
		return 0;
	}

	Process32First( proc_snapshot, &proc_info );
	if ( !wcscmp( proc_info.szExeFile, proc_name.c_str( ) ) ) {
		CloseHandle( proc_snapshot );
		return proc_info.th32ProcessID;
	}

	while ( Process32Next( proc_snapshot, &proc_info ) ) {
		if ( !wcscmp( proc_info.szExeFile, proc_name.c_str( ) ) ) {
			CloseHandle( proc_snapshot );
			return proc_info.th32ProcessID;
		}
	}

	CloseHandle( proc_snapshot );
	return 0;
}

auto efi_hook::efi_t::get_kernel_image( const std::string& module_name ) -> const std::uint64_t {
	void* buffer = nullptr;
	unsigned long buffer_size = 0;

	auto status = NtQuerySystemInformation(
		static_cast< SYSTEM_INFORMATION_CLASS >( SystemModuleInformation ),
		buffer,
		buffer_size,
		&buffer_size
	);

	while ( status == STATUS_INFO_LENGTH_MISMATCH ) {
		VirtualFree( buffer, NULL, MEM_RELEASE );
		buffer = VirtualAlloc( nullptr, buffer_size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE );
		status = NtQuerySystemInformation( static_cast< SYSTEM_INFORMATION_CLASS >( SystemModuleInformation ), buffer, buffer_size, &buffer_size );
	}

	if ( !NT_SUCCESS( status ) ) {
		VirtualFree( buffer, NULL, MEM_RELEASE );
		return 0;
	}

	const auto modules = static_cast< PRTL_PROCESS_MODULES >( buffer );
	for ( auto idx = 0u; idx < modules->NumberOfModules; ++idx ) {
		const auto current_module_name = std::string( reinterpret_cast< char* >( modules->Modules [ idx ].FullPathName ) + modules->Modules [ idx ].OffsetToFileName );
		if ( !_stricmp( current_module_name.c_str( ), module_name.c_str( ) ) ) {
			const auto result = reinterpret_cast< uint64_t >( modules->Modules [ idx ].ImageBase );
			VirtualFree( buffer, NULL, MEM_RELEASE );
			return result;
		}
	}

	VirtualFree( buffer, NULL, MEM_RELEASE );
	return 0;
}

auto efi_hook::efi_t::get_kernel_export( const std::uint64_t image_base, const std::string& module_name ) -> const std::uint64_t {
	const auto ntoskrnl_module = LoadLibraryA( hash_str( "ntoskrnl.exe" ) );
	if ( !ntoskrnl_module ) return 0;

	auto address = reinterpret_cast< std::uint64_t >( GetProcAddress( ntoskrnl_module, module_name.c_str( ) ) );
	if ( !address ) return 0;

	address = address - reinterpret_cast< std::uint64_t >( ntoskrnl_module );
	address = address + image_base;

	FreeLibrary( ntoskrnl_module );

	return address;
}

#endif