#include "scene_hooks.hpp"

#include <windows.h>

#include <algorithm>
#include <atomic>
#include <cstring>
#include <mutex>
#include <unordered_map>

#include "log.hpp"

namespace t5de::scene_hooks {

namespace {

// IMVU's SceneViewer::step() and ModelInstance::updateCalSkeleton(): methods, called with the object in ECX.
using Method = void(__thiscall*)(void* self);

Method originalStep = nullptr;
Method originalUpdateCalSkeleton = nullptr;
std::atomic<std::int64_t> stepTime = 0;
std::atomic<std::uint64_t> steps = 0;
std::atomic<bool> stepHooked = false;
std::atomic<bool> poseHooked = false;

// zero.dll's constants: a tick, in milliseconds (33), and an animation frame, in seconds (1/30).
const unsigned* millisecondsPerFrame = nullptr;
const float* secondsPerFrame = nullptr;

LARGE_INTEGER frequency()
{
    static const LARGE_INTEGER value = [] {
        LARGE_INTEGER f;
        QueryPerformanceFrequency(&f);
        return f;
    }();
    return value;
}

// Our SceneViewer::step: records the tick, then steps as IMVU's does. __fastcall takes ECX like __thiscall (EDX unused).
// Each scene viewer steps in turn at a tick: steps closer than 5 ms are the same tick.
void __fastcall step(void* self, void* /*edx*/)
{
    LARGE_INTEGER now;
    QueryPerformanceCounter(&now);
    if (now.QuadPart - stepTime.load() > frequency().QuadPart / 200) {
        ++steps;
    }
    stepTime = now.QuadPart;
    originalStep(self);
}

// IMVU's objects, as far as posing between ticks reads and writes them (from SceneWindow.dll's ModelInstance::step,
// which copies each effect's time and weight into its Cal3D animation at every tick, and updateCalSkeleton):
// - ModelInstance: +0x2c, a flag: its skeleton must be posed again; +0x120, its animation effects, a std::list (a pointer
//   to the list's head node; a node: next, previous, then the effect).
// - Effect: +0x34, playing (a flag); +0x3c, its time, in seconds; +0x90, its Cal3D animation.
// - The Cal3D animation: +0x8, its time, which the mixer poses the skeleton at (CalMixer::updateSkeleton).
constexpr std::size_t modelPoseDirty = 0x2c;
constexpr std::size_t modelEffects = 0x120;
constexpr std::size_t effectPlaying = 0x34;
constexpr std::size_t effectTime = 0x3c;
constexpr std::size_t effectAnimation = 0x90;
constexpr std::size_t animationTime = 0x8;

template <class T>
T& field(void* object, std::size_t offset)
{
    return *reinterpret_cast<T*>(static_cast<std::byte*>(object) + offset);
}

// Each effect's time at the last two ticks it was posed at, to tell an animation playing on (one frame a tick) from
// one held, looping back or just started (posed as it is).
struct EffectHistory {
    std::uint64_t step = 0;
    float time = 0;
    float previous = -1;
};
std::unordered_map<void*, EffectHistory> histories;

// How far into the current tick the client is, from 0 to just under 1; 0 if unknown.
float tick_fraction()
{
    const std::int64_t last = stepTime.load();
    if (last == 0 || millisecondsPerFrame == nullptr || *millisecondsPerFrame == 0) {
        return 0;
    }
    LARGE_INTEGER now;
    QueryPerformanceCounter(&now);
    const double elapsed = static_cast<double>(now.QuadPart - last) * 1000.0 / static_cast<double>(frequency().QuadPart);
    return static_cast<float>(std::clamp(elapsed / *millisecondsPerFrame, 0.0, 0.999));
}

// Our ModelInstance::updateCalSkeleton, called where ModelInstance::render poses the model before drawing it: between
// ticks, poses it at the time between them. Each playing animation that moved on a frame at the last tick is set
// ahead by the fraction of a frame since, the skeleton posed (by IMVU's), and the times put back, with the skeleton
// marked to be posed again, so everything else sees the tick's pose.
void __fastcall updateCalSkeleton(void* model, void* /*edx*/)
{
    const float fraction = tick_fraction();
    if (fraction < 0.01f || secondsPerFrame == nullptr) {
        originalUpdateCalSkeleton(model);
        return;
    }
    const std::uint64_t tick = steps.load();
    const float frame = *secondsPerFrame;

    struct Saved {
        float* time;
        float value;
    };
    Saved saved[64];
    std::size_t count = 0;
    void* head = field<void*>(model, modelEffects);
    for (void* node = head != nullptr ? field<void*>(head, 0) : nullptr; node != nullptr && node != head && count < 64;
         node = field<void*>(node, 0)) {
        void* effect = field<void*>(node, 8);
        if (effect == nullptr || field<std::uint8_t>(effect, effectPlaying) == 0) {
            continue;
        }
        void* animation = field<void*>(effect, effectAnimation);
        if (animation == nullptr) {
            continue;
        }
        const float time = field<float>(effect, effectTime);
        EffectHistory& history = histories[effect];
        if (history.step != tick) {
            history.previous = history.step + 1 == tick ? history.time : -1;
            history.time = time;
            history.step = tick;
        }
        const float moved = history.time - history.previous;
        if (history.previous < 0 || moved < frame * 0.5f || moved > frame * 1.5f) {
            continue; // held, looped back, or just started: as it is
        }
        float& posedTime = field<float>(animation, animationTime);
        saved[count++] = {&posedTime, posedTime};
        posedTime = time + fraction * frame;
    }
    if (count == 0) {
        originalUpdateCalSkeleton(model);
        return;
    }
    field<std::uint8_t>(model, modelPoseDirty) = 1;
    originalUpdateCalSkeleton(model);
    for (std::size_t i = 0; i < count; ++i) {
        *saved[i].time = saved[i].value;
    }
    field<std::uint8_t>(model, modelPoseDirty) = 1;

    // Forget effects not posed for a while (they may be gone).
    if (histories.size() > 1024) {
        std::erase_if(histories, [tick](const auto& entry) { return entry.second.step + 30 < tick; });
    }
}

// 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;
}

// Redirects the call instruction (E8, then the target relative to the next instruction) at site, if it calls
// expected; returns whether it did.
bool redirect_call(std::byte* site, void* expected, void* replacement)
{
    std::int32_t offset = 0;
    std::memcpy(&offset, site + 1, sizeof(offset));
    if (site[0] != std::byte{0xe8} || site + 5 + offset != expected) {
        return false;
    }
    const auto replacementOffset = static_cast<std::int32_t>(static_cast<std::byte*>(replacement) - (site + 5));
    DWORD protection = 0;
    if (!VirtualProtect(site, 5, PAGE_EXECUTE_READWRITE, &protection)) {
        return false;
    }
    std::memcpy(site + 1, &replacementOffset, sizeof(replacementOffset));
    VirtualProtect(site, 5, protection, &protection);
    FlushInstructionCache(GetCurrentProcess(), site, 5);
    return true;
}

void install_step(HMODULE sceneWindow)
{
    auto* vtable = reinterpret_cast<void**>(GetProcAddress(sceneWindow, "??_7SceneViewer@@6BIFrameListener@@@"));
    auto* stepMethod = 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");
        return;
    }
    originalStep = reinterpret_cast<Method>(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");
        return;
    }
    stepHooked = true;
}

void install_pose(HMODULE sceneWindow)
{
    HMODULE zero = GetModuleHandleW(L"zero.dll");
    millisecondsPerFrame =
        zero == nullptr ? nullptr : reinterpret_cast<const unsigned*>(GetProcAddress(zero, "?millisecondsPerFrame@@3IB"));
    secondsPerFrame = zero == nullptr ? nullptr : reinterpret_cast<const float*>(GetProcAddress(zero, "?secondsPerFrame@@3MB"));
    auto* render = reinterpret_cast<std::byte*>(GetProcAddress(
        sceneWindow, "?render@ModelInstance@@QAEXPAVRenderManager@@ABVCullingFrustum@@ABVTransform@@2ABUFogState@@"
                     "ABULightState@@_NAAVSceneStats@@ABVRenderSettings@@@Z"));
    auto* update = GetProcAddress(sceneWindow, "?updateCalSkeleton@ModelInstance@@AAEXXZ");
    if (millisecondsPerFrame == nullptr || secondsPerFrame == nullptr || render == nullptr || update == nullptr) {
        log::info("scene hooks: ModelInstance::render, updateCalSkeleton or zero.dll's frame lengths not found");
        return;
    }
    // ModelInstance::render poses the model first: its 5th instruction, at +0x12, calls updateCalSkeleton.
    originalUpdateCalSkeleton = reinterpret_cast<Method>(update);
    if (!redirect_call(render + 0x12, reinterpret_cast<void*>(update), reinterpret_cast<void*>(&updateCalSkeleton))) {
        log::info("scene hooks: ModelInstance::render doesn't call updateCalSkeleton where expected");
        return;
    }
    poseHooked = true;
}

} // namespace

void install()
{
    static std::once_flag once;
    std::call_once(once, [] {
        HMODULE sceneWindow = GetModuleHandleW(L"SceneWindow.dll");
        if (sceneWindow == nullptr) {
            log::info("scene hooks: SceneWindow.dll not loaded");
            return;
        }
        install_step(sceneWindow);
        if (stepHooked) {
            install_pose(sceneWindow); // posing between ticks needs the ticks
        }
        log::info("scene hooks: ticks {}, posing between ticks {}", stepHooked ? "recorded" : "not recorded",
                  poseHooked ? "on" : "off");
    });
}

bool installed() { return stepHooked; }

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

} // namespace t5de::scene_hooks