#pragma once

#include <windows.h>

#include <cassert>
#include <cstring>
#include <initializer_list>

namespace rem
{
    void Patch(DWORD address, const void* data, UINT length);
    void PatchBytes(DWORD address, std::initializer_list<BYTE> bytes);
    void Nop(DWORD address, UINT length);
    void ReadWriteProtect(DWORD address, DWORD size);

    template <typename Type>
    inline void Patch(DWORD address, Type value)
    {
        Patch(address, &value, sizeof(Type));
    }

    template <typename Type>
    inline Type& ValueAt(DWORD address)
    {
        ReadWriteProtect(address, sizeof(Type));
        return *reinterpret_cast<Type*>(address);
    }

    template <typename Func>
    inline Func FunctionAt(DWORD address)
    {
        return reinterpret_cast<Func>(address);
    }

    template <typename Func>
    Func SetRelPointer(DWORD location, Func hook)
    {
        DWORD& relative = ValueAt<DWORD>(location);
        DWORD original = location + relative + 4;
        DWORD hookAddress = *reinterpret_cast<DWORD*>(&hook);
        relative = hookAddress - (location + 4);
        return FunctionAt<Func>(original);
    }

    template <typename Func>
    void Hook(DWORD address, Func hook, UINT instructionLength, bool jump = false)
    {
        assert(instructionLength >= 5);
        Patch<BYTE>(address, jump ? 0xE9 : 0xE8);
        SetRelPointer(address + 1, hook);

        if (instructionLength > 5)
            Nop(address + 5, instructionLength - 5);
    }

    template <typename Func>
    Func Trampoline(DWORD address, Func hook, UINT instructionLength)
    {
        BYTE* gateway = static_cast<BYTE*>(
            VirtualAlloc(nullptr, instructionLength + 5, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE));

        ReadWriteProtect(address, instructionLength);
        memcpy(gateway, reinterpret_cast<void*>(address), instructionLength);

        Hook(address, hook, instructionLength, true);
        Hook(reinterpret_cast<DWORD>(gateway + instructionLength), FunctionAt<Func>(address + instructionLength), 5, true);

        return FunctionAt<Func>(reinterpret_cast<DWORD>(gateway));
    }

    template <typename Func>
    void CleanupTrampoline(Func trampoline)
    {
        VirtualFree(reinterpret_cast<void*>(*reinterpret_cast<DWORD*>(&trampoline)), 0, MEM_RELEASE);
    }
}