#include "scene_hooks.hpp"

#include <windows.h>

#include <atomic>
#include <mutex>

#include "log.hpp"

namespace t5de::scene_hooks {

namespace {

// IMVU's SceneViewer::step(), a virtual method: called with the scene viewer in ECX.
using StepFunction = void(__thiscall*)(void* self);

StepFunction originalStep = nullptr;
std::atomic<std::int64_t> stepTime = 0;
std::atomic<bool> hooked = false;

// Our SceneViewer::step: records the tick, then steps as IMVU's does. __fastcall takes ECX like __thiscall (EDX unused).
void __fastcall step(void* self, void* /*edx*/)
{
    LARGE_INTEGER now;
    QueryPerformanceCounter(&now);
    stepTime = now.QuadPart;
    originalStep(self);
}

// Replaces the vtable entry at slot, if it holds expected; returns whether it did.
bool replace(void** slot, void* expected, void* replacement)
{
    if (*slot != expected) {
        return false;
    }
    DWORD protection = 0;
    if (!VirtualProtect(slot, sizeof(void*), PAGE_READWRITE, &protection)) {
        return false;
    }
    *slot = replacement;
    VirtualProtect(slot, sizeof(void*), protection, &protection);
    FlushInstructionCache(GetCurrentProcess(), slot, sizeof(void*));
    return true;
}

} // namespace

void install()
{
    static std::once_flag once;
    std::call_once(once, [] {
        HMODULE sceneWindow = GetModuleHandleW(L"SceneWindow.dll");
        auto* vtable = sceneWindow == nullptr
                           ? nullptr
                           : reinterpret_cast<void**>(GetProcAddress(sceneWindow, "??_7SceneViewer@@6BIFrameListener@@@"));
        auto* stepMethod = sceneWindow == nullptr ? nullptr : GetProcAddress(sceneWindow, "?step@SceneViewer@@UAEXXZ");
        // IFrameListener's vtable: its destructor, then step (what stepFrame calls).
        if (vtable == nullptr || stepMethod == nullptr) {
            log::info("scene hooks: SceneViewer's vtable or step not found; not installed");
            return;
        }
        originalStep = reinterpret_cast<StepFunction>(stepMethod);
        if (!replace(&vtable[1], reinterpret_cast<void*>(stepMethod), reinterpret_cast<void*>(&step))) {
            log::info("scene hooks: SceneViewer's vtable doesn't hold its step; not installed");
            return;
        }
        hooked = true;
        log::info("scene hooks: installed");
    });
}

bool installed() { return hooked; }

std::int64_t lastStep() { return stepTime; }

} // namespace t5de::scene_hooks