#include <iostream>
#include <windows.h>
#include <thread>

#include <fxn/invoker.hpp>

void spawn_car()
{
    auto model = fxn::HashString("adder");
    constexpr auto REQUEST_MODEL = 0xEC9DAA34BBB4658C; // REQUEST_MODEL
    constexpr auto HAS_MODEL_LOADED = 0x6252BC0DD8A320DB; // HAS_MODEL_LOADED
    constexpr auto CREATE_VEHICLE = 0x5779387E956077A6; // CREATE_VEHICLE
    constexpr auto SET_PED_INTO_VEHICLE = 0x73CAFD2038E812B3; // SET_PED_INTO_VEHICLE
    constexpr auto PLAYER_PED_ID = 0x4A8C381C258A124D; // PLAYER_PED_ID
    constexpr auto GET_ENTITY_COORDS = 0xD1A6A821F5AC81DB; // GET_ENTITY_COORDS
    constexpr auto SET_MODEL_AS_NO_LONGER_NEEDED = 0x55098D9E9AD58806; // SET_MODEL_AS_NO_LONGER_NEEDED

    auto& invoker = fxn::NativeInvoker::GetInstance();
    invoker.Schedule<void>(REQUEST_MODEL, model);
    while (!invoker.Schedule<bool>(HAS_MODEL_LOADED, model))
    {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }

    int playerPedId = invoker.Schedule<int>(PLAYER_PED_ID);
    auto playerCoords = invoker.Schedule<fxn::Vector3>(GET_ENTITY_COORDS, playerPedId, true);
    int vehicle = invoker.Schedule<int>(CREATE_VEHICLE, model, playerCoords.x + 5.0f, playerCoords.y, playerCoords.z, 0.0f, true, false, false);
    invoker.Schedule<void>(SET_PED_INTO_VEHICLE, playerPedId, vehicle, -1);
    invoker.Schedule<void>(SET_MODEL_AS_NO_LONGER_NEEDED, model);
}

static void nullsub()
{
    std::printf("[natives_test_dll] Nullsub thread executed.\n");

    std::this_thread::sleep_for(std::chrono::seconds(40));
}

static void* GetInitThreadStub()
{
    static void* addr = []() -> void*
    {   
        void* mem = VirtualAllocEx(GetCurrentProcess(), nullptr, 0x100, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
        if (!mem)
        {
            return nullptr;
        }

        unsigned char stub[] =
        {
            0x40, 0x53,
            0x48, 0x83, 0xec, 0x20,
            0x48, 0xB8,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0xFF, 0xE0, // jmp rax
        };

        auto kernel32 = GetModuleHandleA("kernel32.dll");
        auto baseThreadInitThunk = GetProcAddress(kernel32, "BaseThreadInitThunk");
        if (!baseThreadInitThunk)
        {
            VirtualFreeEx(GetCurrentProcess(), mem, 0, MEM_RELEASE);
            return nullptr;
        }

        *reinterpret_cast<void**>(stub + 8) = reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(baseThreadInitThunk) + 0x6);
        std::memcpy(mem, stub, sizeof(stub));

        return mem;
    }();

    return addr;
}

void CreateThreadBypass()
{
    HANDLE hThread = CreateThread(nullptr, 0, reinterpret_cast<LPTHREAD_START_ROUTINE>(nullsub), nullptr, CREATE_SUSPENDED, nullptr);
    if (!hThread)
    {
        return;
    }

    CONTEXT ctx = { NULL };
    ctx.ContextFlags = CONTEXT_INTEGER; 
    if (!GetThreadContext(hThread, &ctx))
    {
        CloseHandle(hThread);
        return;
    }

    ctx.Rip = reinterpret_cast<DWORD64>(nullsub);

    if (!SetThreadContext(hThread, &ctx))
    {
        CloseHandle(hThread);
        return;
    }

    ResumeThread(hThread);
    CloseHandle(hThread);
}

void HookThreadCreation()
{
    static std::once_flag flag;
    std::call_once(flag, []
    {
        auto ntdll = GetModuleHandleA("ntdll.dll");
        if (!ntdll)
        {
            return;
        }

        auto startUserThreadAddr = GetProcAddress(ntdll, "RtlUserThreadStart");
        if (!startUserThreadAddr)
        {
            return;
        }

        constexpr auto offset = 0x7;
        int32_t rvaOffset = *reinterpret_cast<int32_t*>(reinterpret_cast<uintptr_t>(startUserThreadAddr) + offset + 3);
        std::printf("[natives_test_dll] RtlUserThreadStart RVA offset: 0x%X\n", rvaOffset);

        void** targetAddr = reinterpret_cast<void**>(reinterpret_cast<uintptr_t>(startUserThreadAddr) + offset + 7 + rvaOffset);
        std::printf("[natives_test_dll] RtlUserThreadStart target address: 0x%p\n", targetAddr);

        DWORD oldProtect;
        if (VirtualProtect(targetAddr, sizeof(void*), PAGE_EXECUTE_READWRITE, &oldProtect))
        {
            *targetAddr = GetInitThreadStub();
            VirtualProtect(targetAddr, sizeof(void*), oldProtect, &oldProtect);
        }
    });
}

void mani_()
{
    static auto& invoker = fxn::NativeInvoker::GetInstance();

    // GIVE_WEAPON_TO_PED
    invoker.HookNative(0xB41DEC3AAC1AA107, [](fxn::NativeContext* context, bool& shouldCallOriginal) -> bool
    {
        int32_t pedId = context->GetArgument<int32_t>(0);
        uint32_t weaponHash = context->GetArgument<uint32_t>(1);
        int32_t ammoCount = context->GetArgument<int32_t>(2);
        bool isHidden = context->GetArgument<bool>(3);
        bool forceInHand = context->GetArgument<bool>(4);

        std::cout << "[natives_test_dll] GIVE_WEAPON_TO_PED called with pedId: " << pedId
                  << ", weaponHash: " << std::hex << weaponHash << std::dec
                  << ", ammoCount: " << ammoCount
                  << ", isHidden: " << isHidden
                  << ", forceInHand: " << forceInHand << std::endl;

        shouldCallOriginal = true;
        return true;
    });

    invoker.Initialize();

    std::cout << "[natives_test_dll] mani_ thread started. Press END to exit." << std::endl;

    while (!GetAsyncKeyState(VK_END))
    {
        if (GetAsyncKeyState(VK_F1))
        {
            fxn::NativeContext context { 0x4A8C381C258A124D }; // PLAYER_PED_ID
            invoker.Schedule(&context);

            int playerPedId = context.GetResult<int>();
            std::cout << "[natives_test_dll] PLAYER_PED_ID: " << playerPedId << std::endl;
        }

        if (GetAsyncKeyState(VK_F2))
        {
            std::cout << "[natives_test_dll] Spawning car..." << std::endl;
            spawn_car();
            std::cout << "[natives_test_dll] Car spawned." << std::endl;
        }

        if (GetAsyncKeyState(VK_F3))
        {
            std::once_flag flag;
            std::call_once(flag, []
            {
                HookThreadCreation();
            });

            CreateThreadBypass();
            std::cout << "[natives_test_dll] Created thread with bypass." << std::endl;
        }

        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }

    std::cout << "[natives_test_dll] Exiting mani_ thread." << std::endl;
}


DWORD WINAPI ThreadProc(LPVOID lpParameter)
{
    Sleep(1000);

    AllocConsole();

    FILE* fDummy;
    freopen_s(&fDummy, "CONOUT$", "w", stdout);

    mani_();

    fclose(fDummy);
    FreeConsole();

    FreeLibraryAndExitThread(static_cast<HMODULE>(lpParameter), 0);
    return 0;
}

void Initialize(LPVOID lpParameter)
{
    AllocConsole();
    FILE* fDummy;

    freopen_s(&fDummy, "CONOUT$", "w", stdout);
    std::cout << "[natives_test_dll] DLL injected successfully." << std::endl;
}

HMODULE g_module = nullptr;

BOOL APIENTRY DllMain(HMODULE hModule, DWORD ulReason, LPVOID lpReserved)
{
    g_module = hModule;

    if (ulReason == DLL_PROCESS_ATTACH)
    {
        DisableThreadLibraryCalls(hModule);

        // fxn::NativeInvoker::GetInstance().Initialize();
        
        HookThreadCreation();

        HANDLE hThread = CreateThread(nullptr, 0, ThreadProc, hModule, 0, nullptr);
        if (hThread)
        {
            CloseHandle(hThread);
        }
    }

    return TRUE;
}