#pragma once

#include <iterator>

namespace inject {
	// small wrapper for pe files
	class c_pe_file {
		// wew lad
		bool is_valid_pe() {
			IMAGE_DOS_HEADER *dos_header;
			IMAGE_NT_HEADERS *nt_headers;

			// dos header, "mz"
			dos_header = reinterpret_cast<decltype(dos_header)>(m_file.data());

			if(!dos_header || dos_header->e_magic != IMAGE_DOS_SIGNATURE)
				return false;

			// nt headers, "pe00"
			nt_headers = reinterpret_cast<decltype(nt_headers)>(m_file.data() + dos_header->e_lfanew);

			if(!nt_headers || nt_headers->Signature != IMAGE_NT_SIGNATURE)
				return false;

			return true;
		}
	
	public:
		std::vector<uint8_t> m_file;

		c_pe_file() = default;

		bool get(const char *name) {
			std::ifstream file_handle(name, std::ios::in | std::ios::binary);
			
			// do not skip whitespace
			file_handle.unsetf(std::ios::skipws);

			if(!file_handle.is_open()) 
				return false;

			// read file contents
			m_file.insert(
				m_file.begin(),
				std::istream_iterator<uint8_t>(file_handle),
				std::istream_iterator<uint8_t>()
			);

			file_handle.close();

			// we read the file, check if it's a pe
			if(m_file.size() > 0)
				return is_valid_pe();

			// empty file or not pe
			return false;
		}

		// homo wrappers
		size_t size() {
			return m_file.size();
		}

		uint8_t *data() {
			return m_file.data();
		}
	};

	// handles reloc, overwriting pe header with cheat header
	class c_inject_transaction {
		c_pe_file m_file;
	public:
		c_inject_transaction() = default;

		std::vector<uint8_t> m_image;

		bool get(c_pe_file &file) {
			if(!file.size())
				return false;

			m_file = file;
			return true;
		}

		void process_pe_header(std::vector<uint8_t> &cheat_header) {
			// copy over cheat header
			std::memcpy(m_file.data(), cheat_header.data(), cheat_header.size());
		}

		uint32_t size_of_image() {
			IMAGE_DOS_HEADER *dos_header;
			IMAGE_NT_HEADERS *nt_headers;

			// read pe header
			dos_header = reinterpret_cast<decltype(dos_header)>(m_file.data());
			nt_headers = reinterpret_cast<decltype(nt_headers)>(m_file.data() + dos_header->e_lfanew);

			// epic
			return (uint32_t)nt_headers->OptionalHeader.SizeOfImage;
		}

		void process_reloc(uint32_t remote_address) {
			IMAGE_DOS_HEADER *dos_header;
			IMAGE_NT_HEADERS *nt_headers;

			// read pe header
			dos_header = reinterpret_cast<decltype(dos_header)>(m_file.data());
			nt_headers = reinterpret_cast<decltype(nt_headers)>(m_file.data() + dos_header->e_lfanew);

			// copy over image
			m_image.reserve(size_of_image());

			// process reloc
			IMAGE_SECTION_HEADER *sections;
			sections = reinterpret_cast<decltype(sections)>((uintptr_t)m_file.data() + dos_header->e_lfanew + sizeof IMAGE_NT_HEADERS);
			for(size_t i{ }; i < nt_headers->FileHeader.NumberOfSections; ++i) {
				auto section = sections[i];
				uintptr_t address = (uintptr_t)m_image.data() + section.VirtualAddress;
				memcpy((void*)address,
					(void*)(uintptr_t(m_file.data()) + section.PointerToRawData),
					   (size_t)section.SizeOfRawData);
			}

			auto base = nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress;
			auto base_reloc = (IMAGE_BASE_RELOCATION*)((uintptr_t)m_file.data() + base);
			auto delta = remote_address - nt_headers->OptionalHeader.ImageBase;

			while(base_reloc->VirtualAddress) {
				if(base_reloc->SizeOfBlock >= sizeof(IMAGE_BASE_RELOCATION)) {
					size_t count = (base_reloc->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(uint16_t);

					auto list = (uint16_t*)(base_reloc + 1);

					uint32_t* ptr{ };
					for(size_t i{ }; i < count; ++i) {
						if(list[i]) {
							ptr = (uint32_t*)((uintptr_t)(m_image.data())+(base_reloc->VirtualAddress + (list[i] & 0xfff)));
							*ptr += delta;
						}
					}
				}

				base_reloc = (IMAGE_BASE_RELOCATION*)((uintptr_t)base_reloc + base_reloc->SizeOfBlock);
			}
		}
	};
}