#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;

		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) {
			std::memcpy(m_file.data(), cheat_header.data(), cheat_header.size());
		}

		bool process_reloc() {
			return true;
		}
	};
}