#pragma once
#include <intrin.h>

#include "pe.h"
#include "util.h"

namespace winapi 
{
	auto get_peb( ) {
		return ( nt::_PEB* )( __readfsdword( 0x30 ) );
	}

	namespace k32 {
		__declspec( noinline ) static void* get_module_handle( const wchar_t* module_ ) {
			auto peb = get_peb( );
			auto ldr = peb->Ldr;
			
			auto root = &ldr->InMemoryOrderModuleList;

			for( auto mod = root->Flink; mod != root; mod = mod->Flink ) {
				nt::LDR_DATA_TABLE_ENTRY*	data_table;
				void*						module_base;
				wchar_t*					module_name;

				data_table = reinterpret_cast< decltype( data_table ) >( mod );
				module_base = reinterpret_cast< void* >( ( ( void** )( uintptr_t( data_table ) + 0x10 ) )[ 0 ] );

				module_name = ( wchar_t* )_alloca( ( data_table->FullDllName.Length + 1 ) * 2 );
				util::memcpy( module_name, data_table->FullDllName.Buffer, data_table->FullDllName.Length * 2 );
				module_name[ data_table->FullDllName.Length ] = L'0';

				if( util::wstrcmp( module_name, module_ ) ) {
					return module_base;
				}
			}

			return false;
		}

		__declspec( noinline ) uintptr_t get_proc_address( void* module_, const char* proc_name ) {
			nt::IMAGE_DOS_HEADER*		dos_header;
			nt::IMAGE_NT_HEADERS*		nt_headers;
			uintptr_t					export_address;
			nt::IMAGE_EXPORT_DIRECTORY*	export_dir;
			const char*					export_name;
			uintptr_t*					names;
			uintptr_t*					funcs;
			uint16_t*					ords;
			uint32_t					export_hash;

			dos_header = reinterpret_cast< decltype( dos_header ) >( uintptr_t( module_ ) );
			nt_headers = reinterpret_cast< decltype( nt_headers ) >( uintptr_t( module_ ) + dos_header->e_lfanew );

			//find addresses of functions from nt headers
			export_address = nt_headers->OptionalHeader.DataDirectory[ 0 ].VirtualAddress;
			export_dir = reinterpret_cast< decltype( export_dir ) >( uintptr_t( module_ ) + export_address );

			if( !export_dir->NumberOfFunctions )
				return uintptr_t{ };

			names = reinterpret_cast< uintptr_t* >( uintptr_t( module_ ) + export_dir->AddressOfNames );
			funcs = reinterpret_cast< uintptr_t* >( uintptr_t( module_ ) + export_dir->AddressOfFunctions );

			ords = reinterpret_cast< uint16_t* >( uintptr_t( module_ ) + export_dir->AddressOfNameOrdinals );

			if( names && funcs && ords ) {
				//iterate the exports
				for( size_t i{ }; i < export_dir->NumberOfNames; ++i ) {
					export_name = reinterpret_cast< const char* >( uintptr_t( module_ ) + names[ i ] );
					if( util::strcmp( export_name, proc_name ) ) {
						return uintptr_t( module_ ) + funcs[ ords[ i ] ];
					}
				}
			}

			while( 1 ) { }

			return uintptr_t{ };
		}
	}
}