#pragma once

namespace PolyHook_Hooks {
	uint64_t hookSmall(void);
	uint64_t hookBranch(uint64_t);
	uint64_t hookRip_relative(void);
	void hookAVX(float num, void* res);
	uint32_t hookRDRAND(void);
	uint32_t hookLoop(uint32_t x);
	uint32_t hookTail_recursion(uint32_t x);
};

class PolyHook : public AbstractHookEngine {
private:
	std::shared_ptr<PLH::Detour> detSmall,
		detBranch,
		detRIPRelative,
		detAVX,
		detRDRAND,
		detLoop,
		detTailRecursion;

	template <typename T>
	bool hook(std::shared_ptr<PLH::Detour> det, void* func, T& original, void* hook) {
		det->SetupHook((BYTE*)func, (BYTE*)hook);
		if (!det->Hook()) {
			return false;
		}

		return (original = det->GetOriginal<T>());
	}

public:
	bool hook_all();
	bool unhook_all();

	PolyHook() : AbstractHookEngine("PolyHook"),
		detSmall(new PLH::Detour),
		detRIPRelative(new PLH::Detour),
		detBranch(new PLH::Detour),
		detAVX(new PLH::Detour),
		detRDRAND(new PLH::Detour),
		detLoop(new PLH::Detour),
		detTailRecursion(new PLH::Detour)
	{
		assert(detSmall->GetType() == PLH::HookType::Detour);
		assert(detBranch->GetType() == PLH::HookType::Detour);
		assert(detRIPRelative->GetType() == PLH::HookType::Detour);
		assert(detAVX->GetType() == PLH::HookType::Detour);
		assert(detRDRAND->GetType() == PLH::HookType::Detour);
		assert(detLoop->GetType() == PLH::HookType::Detour);
		assert(detTailRecursion->GetType() == PLH::HookType::Detour);
	}

	friend uint64_t PolyHook_Hooks::hookSmall(void);
	friend uint64_t PolyHook_Hooks::hookBranch(uint64_t);
	friend uint64_t PolyHook_Hooks::hookRip_relative(void);

	friend void PolyHook_Hooks::hookAVX(float num, void* res);
	friend uint32_t PolyHook_Hooks::hookRDRAND(void);
	friend uint32_t PolyHook_Hooks::hookLoop(uint32_t num);
	friend uint32_t PolyHook_Hooks::hookTail_recursion(uint32_t x);
};