#ifndef PDB_CPP
#define PDB_CPP

#include <core/inc.hpp>
#include <string>
#include <fstream>
#include <libloaderapi.h>
#include <cstdint>
#include <rpcasync.h>
#include <wtypes.h>
#include <TlHelp32.h>
#include <vector>
#include <filesystem>
#include <combaseapi.h>
#include <DbgHelp.h>
#pragma comment(lib, "DbgHelp.lib")
#pragma comment(lib, "Urlmon.lib")

auto pdb_interface::c_pdb::download_pdb_to_temp( std::string service ) -> std::string {

	auto ntos_executable_path = std::string( std::getenv( hash_str( "systemroot" ) ) ) + service;

	auto pdb_download_path =    std::string( hash_str( "C:\\Windows\\Temp" ) );

	if ( pdb_download_path[ pdb_download_path.size( ) - 1 ] != '\\' )
		pdb_download_path += hash_str("\\");

	std::ifstream file_object_handle( ntos_executable_path, std::ios::binary | std::ios::ate );

	std::streamsize file_size = file_object_handle.tellg( );

	file_object_handle.seekg( 0, std::ios::beg );

	std::vector<char> file_buffer( file_size );

	if ( !file_object_handle.read( file_buffer.data( ), file_size ) || file_size == 0 ) {
		SetLastError( ERROR_ACCESS_DENIED );
		return hash_str( "NONE" );
	}

	auto pdb_path = pdb_download_path + random_string( 16 ) + hash_str(".pdb");

	auto dos_header = ( IMAGE_DOS_HEADER* ) file_buffer.data( );
	auto nt_header = ( IMAGE_NT_HEADERS* ) ( file_buffer.data( ) + dos_header->e_lfanew );
	auto file_header = &nt_header->FileHeader;

	auto optional_file_header64 = ( IMAGE_OPTIONAL_HEADER64* ) NULL;
	auto optional_file_header32 = ( IMAGE_OPTIONAL_HEADER32* ) NULL;

	auto x86 = FALSE;

	if ( file_header->Machine == IMAGE_FILE_MACHINE_AMD64 ) {
		optional_file_header64 = ( IMAGE_OPTIONAL_HEADER64* ) ( &nt_header->OptionalHeader );
	}
	else if ( file_header->Machine == IMAGE_FILE_MACHINE_I386 ) {
		optional_file_header32 = ( IMAGE_OPTIONAL_HEADER32* ) ( &nt_header->OptionalHeader );
		x86 = TRUE;
	}
	else {
		SetLastError( ERROR_NOT_SUPPORTED );
		return hash_str( "NONE" );
	}

	auto image_size = x86 ? optional_file_header32->SizeOfImage : optional_file_header64->SizeOfImage;
	auto image_buffer = ( PBYTE ) malloc( image_size );
	
	if ( !image_buffer ) {
		SetLastError( ERROR_NOT_ENOUGH_MEMORY );
		return hash_str( "NONE" );
	}

	memcpy( image_buffer, file_buffer.data( ), x86 ? optional_file_header32->SizeOfHeaders : optional_file_header64->SizeOfHeaders );
	
	IMAGE_SECTION_HEADER* current_section_header = IMAGE_FIRST_SECTION( nt_header );

	for ( UINT i = 0; i != file_header->NumberOfSections; ++i, ++current_section_header )
	{
		if ( current_section_header->SizeOfRawData )
		{
			memcpy( image_buffer + current_section_header->VirtualAddress, file_buffer.data( ) + current_section_header->PointerToRawData, current_section_header->SizeOfRawData );
		}
	}

	IMAGE_DATA_DIRECTORY* data_directory = nullptr;
	if ( x86 )
	{
		data_directory = &optional_file_header32->DataDirectory[ IMAGE_DIRECTORY_ENTRY_DEBUG ];
	}
	else
	{
		data_directory = &optional_file_header64->DataDirectory[ IMAGE_DIRECTORY_ENTRY_DEBUG ];
	}

	auto image_debug_directory = ( IMAGE_DEBUG_DIRECTORY* ) ( image_buffer + data_directory->VirtualAddress );
	if ( !data_directory->Size || IMAGE_DEBUG_TYPE_CODEVIEW != image_debug_directory->Type ) 
	{
		free( image_buffer );
		SetLastError( ERROR_NOT_SUPPORTED );
		return hash_str( "NONE" );
	}

	auto image_pdb_info = ( pdb_info* ) ( image_buffer + image_debug_directory->AddressOfRawData );
	if ( image_pdb_info->signature != 0x53445352 )
	{
		free( image_buffer );
		SetLastError( ERROR_NOT_SUPPORTED );
		return hash_str( "NONE" );
	}

	if ( !std::filesystem::exists( pdb_path ) )
	{
		wchar_t w_GUID[ 100 ] = { 0 };
		if ( !StringFromGUID2( image_pdb_info->guid, w_GUID, 100 ) )
		{
			free( image_buffer );
			SetLastError( ERROR_NOT_SUPPORTED );
			return hash_str( "NONE" );
		}

		char a_GUID[ 100 ]{ 0 };
		size_t l_GUID = 0;
		if ( wcstombs_s( &l_GUID, a_GUID, w_GUID, sizeof( a_GUID ) ) || !l_GUID )
		{
			free( image_buffer );
			SetLastError( ERROR_NOT_SUPPORTED );
			return hash_str( "NONE" );
		}

		char guid_filtered[ 256 ] = { 0 };
		for ( UINT i = 0; i != l_GUID; ++i )
		{
			if ( ( a_GUID[ i ] >= '0' && a_GUID[ i ] <= '9' ) || ( a_GUID[ i ] >= 'A' && a_GUID[ i ] <= 'F' ) || ( a_GUID[ i ] >= 'a' && a_GUID[ i ] <= 'f' ) )
			{
				guid_filtered[ strlen( guid_filtered ) ] = a_GUID[ i ];
			}
		}

		char age[ 3 ] = { 0 };
		_itoa_s( image_pdb_info->age, age, 10 );

		std::string url = hash_str( "https://msdl.microsoft.com/download/symbols/" );
		url += image_pdb_info->pdb_file_name;
		url += hash_str("/");
		url += guid_filtered;
		url += age;
		url += hash_str("/");
		url += image_pdb_info->pdb_file_name;

		HRESULT hr = URLDownloadToFileA( NULL, url.c_str( ), pdb_path.c_str( ), NULL, NULL );
		if ( FAILED( hr ) )
		{
			free( image_buffer );
			return hash_str( "NONE" );
		}

		free( image_buffer );
	}

	return pdb_path;
}

auto pdb_interface::c_pdb::load_pdb( 
	std::string pdb_path,
	pdb_context pdb_object ) -> bool {

	memset( pdb_object, 0, sizeof( pdb_context ) );

	WIN32_FILE_ATTRIBUTE_DATA file_attribution_data{ 0 };

	if ( !GetFileAttributesExA( pdb_path.c_str( ), GetFileExInfoStandard, &file_attribution_data ) ) {
		return false;
	}

	auto pdb_size = file_attribution_data.nFileSizeLow;
	auto pdb_file_handle = CreateFileA( pdb_path.c_str( ), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, NULL, NULL );
	if ( pdb_file_handle == INVALID_HANDLE_VALUE ) {
		return false;
	}

	auto process_handle = OpenProcess( PROCESS_QUERY_LIMITED_INFORMATION, FALSE, GetCurrentProcessId( ) );
	if ( !process_handle ) {
		CloseHandle( pdb_file_handle );
		return false;
	}

	if ( !SymInitialize( process_handle, pdb_path.c_str( ), FALSE ) )
	{
		CloseHandle( process_handle );
		CloseHandle( pdb_file_handle );
		return false;
	}

	SymSetOptions( SYMOPT_UNDNAME | SYMOPT_DEFERRED_LOADS | SYMOPT_AUTO_PUBLICS | SYMOPT_DEBUG | SYMOPT_LOAD_ANYTHING );

	auto symbol_table = SymLoadModuleEx( process_handle, NULL, pdb_path.c_str( ), NULL, EZ_PDB_BASE_OF_DLL, pdb_size, NULL, NULL );
	if ( !symbol_table )
	{
		SymCleanup( process_handle );
		CloseHandle( process_handle );
		CloseHandle( pdb_file_handle );
		return false;
	}

	pdb_object->pdb_file = pdb_file_handle;
	pdb_object->process = process_handle;

	return true;
}

auto pdb_interface::c_pdb::get_rva( 
	pdb_context pdb_object, 
	std::string symbol_name ) -> std::uint64_t {

	SYMBOL_INFO symbol_information = { 0 };

	symbol_information.SizeOfStruct = sizeof( SYMBOL_INFO );

	if ( !SymFromName( pdb_object->process, symbol_name.c_str( ), &symbol_information ) ) {
		return ( std::uint64_t ) -1;
	}

	return ( std::uint64_t )( symbol_information.Address - symbol_information.ModBase );
}

auto pdb_interface::c_pdb::get_struct_property_offset(
	pdb_context pdb_object,
	std::string struct_name,
	std::wstring property_name ) -> std::uint64_t {

	auto symbol_info_size = (ULONG) sizeof( SYMBOL_INFO ) + MAX_SYM_NAME * sizeof( TCHAR );
	auto symbol_info = ( SYMBOL_INFO* ) malloc( symbol_info_size );
	if ( !symbol_info ) {
		return ( std::uint64_t ) -1;
	}

	ZeroMemory( symbol_info, symbol_info_size );
	symbol_info->SizeOfStruct = sizeof( SYMBOL_INFO );
	symbol_info->MaxNameLen = MAX_SYM_NAME;

	if ( !SymGetTypeFromName( pdb_object->process, EZ_PDB_BASE_OF_DLL, struct_name.c_str( ), symbol_info ) ) {
		return ( std::uint64_t ) -1;
	}

	TI_FINDCHILDREN_PARAMS find_children_parameters = { 0 };
	if ( !SymGetTypeInfo( pdb_object->process, EZ_PDB_BASE_OF_DLL, symbol_info->TypeIndex, TI_GET_CHILDRENCOUNT, &find_children_parameters ) ) {
		free( symbol_info );
		return ( std::uint64_t ) -1;
	}

	auto children_parameters_size = ( ULONG ) sizeof( TI_FINDCHILDREN_PARAMS ) + find_children_parameters.Count * sizeof( ULONG );
	auto children_parameters = ( TI_FINDCHILDREN_PARAMS* ) malloc( children_parameters_size );
	if ( !children_parameters ) {
		free( symbol_info );
		return ( std::uint64_t ) -1;
	}

	ZeroMemory( children_parameters, children_parameters_size );
	memcpy( children_parameters, &find_children_parameters, sizeof( TI_FINDCHILDREN_PARAMS ) );
	
	if ( !SymGetTypeInfo( pdb_object->process, EZ_PDB_BASE_OF_DLL, symbol_info->TypeIndex, TI_FINDCHILDREN, children_parameters ) ) {
		free( children_parameters );
		free( symbol_info );
		return ( std::uint64_t ) -1;
	}

	for ( ULONG i = children_parameters->Start; i < children_parameters->Count; i++ )
	{
		WCHAR* pSymName = NULL;
		std::uint64_t offset = 0;
		if ( !SymGetTypeInfo( pdb_object->process, EZ_PDB_BASE_OF_DLL, children_parameters->ChildId[ i ], TI_GET_OFFSET, &offset ) )
		{
			free( children_parameters );
			free( symbol_info );
			return ( std::uint64_t ) -1;
		}
		if ( !SymGetTypeInfo( pdb_object->process, EZ_PDB_BASE_OF_DLL, children_parameters->ChildId[ i ], TI_GET_SYMNAME, &pSymName ) )
		{
			free( children_parameters );
			free( symbol_info );
			return ( std::uint64_t ) -1;
		}
		if ( pSymName )
		{
			if ( wcscmp( pSymName, property_name.c_str( ) ) == 0 )
			{
				LocalFree( pSymName );
				free( children_parameters );
				free( symbol_info );
				return offset;
			}
		}
	}

	free( children_parameters );
	free( symbol_info );
	return ( std::uint64_t ) -1;
}

auto pdb_interface::c_pdb::get_struct_size(
	pdb_context pdb_object,
	std::string struct_name ) -> std::uint64_t {

	auto symbol_info_size = sizeof( SYMBOL_INFO ) + MAX_SYM_NAME * sizeof( TCHAR );
	auto symbol_info = ( SYMBOL_INFO* ) malloc( symbol_info_size );
	if ( !symbol_info ) {
		return ( std::uint64_t ) -1;
	}

	ZeroMemory( symbol_info, symbol_info_size );
	
	symbol_info->SizeOfStruct = sizeof( SYMBOL_INFO );
	symbol_info->MaxNameLen = MAX_SYM_NAME;

	if ( !SymGetTypeFromName( pdb_object->process, EZ_PDB_BASE_OF_DLL, struct_name.c_str( ), symbol_info ) ) {
		return ( std::uint64_t ) -1;
	}

	return symbol_info->Size;
}

auto pdb_interface::c_pdb::unload_pdb( pdb_context pdb_object ) -> void {

	SymUnloadModule64( pdb_object->process, EZ_PDB_BASE_OF_DLL );
	SymCleanup( pdb_object->process );
	CloseHandle( pdb_object->process );
	CloseHandle( pdb_object->pdb_file );
}

#endif // ! guard