#include "scene_hooks.hpp"

#include <windows.h>

#include <algorithm>
#include <atomic>
#include <cmath>
#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;
// IMVU's CameraController::getCameraTransform(): returns a Transform by value (constructed in result).
using GetTransform = void*(__thiscall*)(const void* self, void* result);
GetTransform originalGetCameraTransform = nullptr;
float cameraEasing = 0; // the share of the way to its target the camera moves at each tick
std::atomic<std::int64_t> stepTime = 0;
std::atomic<std::uint64_t> steps = 0;
// How many ticks the last step covered: the client steps a scene several ticks at once when it falls behind.
std::atomic<int> stepTicks = 1;
std::atomic<bool> stepHooked = false;
std::atomic<bool> poseHooked = false;
std::atomic<bool> cameraHooked = false;
std::atomic<bool> particlesHooked = false;

// IMVU's ParticleSystem::render(RenderManager*, const Transform&, const Transform&): a method with 3 arguments.
using RenderParticles = void(__thiscall*)(void* self, void* renderManager, const void* a, const void* b);
RenderParticles originalRenderParticles = nullptr;
// The particle library's (ParticleLib.dll: PAPI, the Particle System API) PContextParticleGroup_t::GetParticlePointer:
// the current group's particles, as floats: returns their count, and sets the first's address, the stride between
// them, and the offsets of their attributes (in floats): position, previous position, size, velocity, previous
// velocity, color, alpha, age, up, rotational velocity, previous up, mass, data.
using GetParticlePointer = unsigned(__thiscall*)(void* group, float*& particles, unsigned& stride, unsigned& position,
                                                 unsigned& positionB, unsigned& size, unsigned& velocity,
                                                 unsigned& velocityB, unsigned& color, unsigned& alpha, unsigned& age,
                                                 unsigned& up, unsigned& rotationalVelocity, unsigned& upB,
                                                 unsigned& mass, unsigned& data);
GetParticlePointer getParticlePointer = nullptr;
// PContextParticleGroup_t::GetGroupCount: how many particle groups there are (GetParticlePointer throws without one).
using GetGroupCount = unsigned(__thiscall*)(void* group);
GetGroupCount getGroupCount = nullptr;

// 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);
    const std::int64_t last = stepTime.load();
    if (now.QuadPart - last > frequency().QuadPart / 200) {
        ++steps;
        if (last != 0 && millisecondsPerFrame != nullptr && *millisecondsPerFrame != 0) {
            const double ticks = static_cast<double>(now.QuadPart - last) * 1000.0 /
                                 static_cast<double>(frequency().QuadPart) / *millisecondsPerFrame;
            stepTicks = std::clamp(static_cast<int>(ticks + 0.5), 1, 5);
        }
        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 (at any rate: the time it
// moved per tick at the last step) 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 at the last step is set ahead by the
// fraction of a tick since, at the rate it moved (up to 4 frames a tick), 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 perTick = (history.time - history.previous) / static_cast<float>(stepTicks.load());
        if (history.previous < 0 || perTick <= frame * 0.01f || perTick > frame * 4) {
            continue; // held, looped back, or just started: as it is
        }
        float& posedTime = field<float>(animation, animationTime);
        saved[count++] = {&posedTime, posedTime};
        posedTime = time + fraction * perTick;
    }
    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; });
    }
}

// The camera (SceneWindow.dll's CameraController::step, at each tick): each of these values moves a share k of the way
// to its target (k is 0.4): the camera's position (+0x118, 3 floats, to +0xf0; panning moves the target), and its
// orbit angles and zoom, in degrees (+0x140, 3 floats, to +0x134; the difference taken between -180 and 180).
// A fraction f of a tick on, the same easing has moved it 1 - (1 - k)^f of the way: exactly, nothing guessed. The point
// the camera looks at (+0x160, 3 floats) is the position plus an offset (setCameraFocus, at each tick), so it moves with
// the position.
constexpr std::size_t cameraPosition = 0x118;
constexpr std::size_t cameraPositionTarget = 0xf0;
constexpr std::size_t cameraAngles = 0x140;
constexpr std::size_t cameraAnglesTarget = 0x134;
constexpr std::size_t cameraFocus = 0x160;
// Nonzero: the camera shows a fixed view ("take"), not eased.
constexpr std::size_t cameraTake = 0xec;

float degrees180(float degrees)
{
    float wrapped = std::fmod(degrees + 180.0f, 360.0f);
    if (wrapped < 0) {
        wrapped += 360.0f;
    }
    return wrapped - 180.0f;
}

// Our CameraController::getCameraTransform: between ticks, the camera where its easing has brought it since the tick.
// The eased values are set ahead for IMVU's code to make the transform from, then put back. __fastcall takes ECX like
// __thiscall, and like it the callee pops the stack argument.
void* __fastcall getCameraTransform(void* camera, void* /*edx*/, void* result)
{
    const float fraction = tick_fraction();
    if (fraction < 0.01f || field<std::int32_t>(camera, cameraTake) != 0) {
        return originalGetCameraTransform(camera, result);
    }
    const float moved = 1.0f - std::pow(1.0f - cameraEasing, fraction);
    float saved[9];
    for (int i = 0; i < 3; ++i) {
        float& position = field<float>(camera, cameraPosition + i * 4);
        float& focus = field<float>(camera, cameraFocus + i * 4);
        saved[i] = position;
        saved[6 + i] = focus;
        const float step = (field<float>(camera, cameraPositionTarget + i * 4) - position) * moved;
        position += step;
        focus += step;
        float& angle = field<float>(camera, cameraAngles + i * 4);
        saved[3 + i] = angle;
        angle = degrees180(angle + degrees180(field<float>(camera, cameraAnglesTarget + i * 4) - angle) * moved);
    }
    void* transform = originalGetCameraTransform(camera, result);
    for (int i = 0; i < 3; ++i) {
        field<float>(camera, cameraPosition + i * 4) = saved[i];
        field<float>(camera, cameraAngles + i * 4) = saved[3 + i];
        field<float>(camera, cameraFocus + i * 4) = saved[6 + i];
    }
    return transform;
}

// Particles (SceneWindow.dll's ParticleSystem::step, at each tick) move by their velocity times the tick's time
// (PContextActionList_t::TimeStep, millisecondsPerFrame / 1000 seconds a tick, then Move). Between ticks, they're drawn
// moved on by their velocity times the time since the tick, then moved back the same way. (ParticleSystem::render
// sorts the particles by distance before drawing them, reordering them in memory: each is moved back by its own
// velocity, wherever it now is.) A ParticleSystem's particle group is at +0x28 (what ParticleSystem::render reads
// them from).
constexpr std::size_t particleGroup = 0x28;

void __fastcall renderParticles(void* system, void* /*edx*/, void* renderManager, const void* a, const void* b)
{
    const float fraction = tick_fraction();
    void* group = static_cast<std::byte*>(system) + particleGroup;
    float* particles = nullptr;
    unsigned stride = 0, position = 0, positionB = 0, size = 0, velocity = 0, velocityB = 0, color = 0, alpha = 0,
             age = 0, up = 0, rotationalVelocity = 0, upB = 0, mass = 0, data = 0;
    // As ParticleSystem::render does: no particles without a group. The library reports errors by throwing; then the
    // particles are drawn as they are.
    const auto particle_pointer = [&] {
        try {
            if (getGroupCount(group) != 0) {
                return getParticlePointer(group, particles, stride, position, positionB, size, velocity, velocityB,
                                          color, alpha, age, up, rotationalVelocity, upB, mass, data);
            }
        } catch (...) {
        }
        return 0u;
    };
    const unsigned count = fraction >= 0.01f ? particle_pointer() : 0;
    if (count == 0 || particles == nullptr || stride == 0) {
        originalRenderParticles(system, renderManager, a, b);
        return;
    }
    const float seconds = fraction * static_cast<float>(*millisecondsPerFrame) / 1000.0f;
    const auto move = [&](float by) {
        for (unsigned i = 0; i < count; ++i) {
            float* particle = particles + std::size_t{i} * stride;
            for (unsigned axis = 0; axis < 3; ++axis) {
                particle[position + axis] += particle[velocity + axis] * by;
            }
        }
    };
    move(seconds);
    originalRenderParticles(system, renderManager, a, b);
    if (particle_pointer() == count && particles != nullptr) { // the same particles, maybe reordered
        move(-seconds);
    }
}

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

void install_particles(HMODULE sceneWindow)
{
    HMODULE particleLib = GetModuleHandleW(L"ParticleLib.dll");
    getParticlePointer = particleLib == nullptr
                             ? nullptr
                             : reinterpret_cast<GetParticlePointer>(GetProcAddress(
                                   particleLib, "?GetParticlePointer@PContextParticleGroup_t@PAPI@@QAEIAAPAMAAI1111111111111@Z"));
    getGroupCount = particleLib == nullptr ? nullptr
                                           : reinterpret_cast<GetGroupCount>(GetProcAddress(
                                                 particleLib, "?GetGroupCount@PContextParticleGroup_t@PAPI@@QAEIXZ"));
    auto* modelRender = reinterpret_cast<std::byte*>(GetProcAddress(
        sceneWindow, "?render@ModelInstance@@QAEXPAVRenderManager@@ABVCullingFrustum@@ABVTransform@@2ABUFogState@@"
                     "ABULightState@@_NAAVSceneStats@@ABVRenderSettings@@@Z"));
    auto* render = GetProcAddress(sceneWindow, "?render@ParticleSystem@@QAEXPAVRenderManager@@ABVTransform@@1@Z");
    if (getParticlePointer == nullptr || getGroupCount == nullptr || modelRender == nullptr || render == nullptr
        || millisecondsPerFrame == nullptr) {
        log::info("scene hooks: ParticleLib's GetParticlePointer, ModelInstance::render or ParticleSystem::render not found");
        return;
    }
    // ModelInstance::render draws its particle systems with a call at +0x114.
    originalRenderParticles = reinterpret_cast<RenderParticles>(render);
    if (!redirect_call(modelRender + 0x114, reinterpret_cast<void*>(render), reinterpret_cast<void*>(&renderParticles))) {
        log::info("scene hooks: ModelInstance::render doesn't call ParticleSystem::render where expected");
        return;
    }
    particlesHooked = true;
}

void install_camera(HMODULE sceneWindow)
{
    auto* vtable = reinterpret_cast<void**>(GetProcAddress(sceneWindow, "??_7CameraController@@6B@"));
    auto* get = GetProcAddress(sceneWindow, "?getCameraTransform@CameraController@@UBE?AVTransform@@XZ");
    auto* step = reinterpret_cast<const std::byte*>(GetProcAddress(sceneWindow, "?step@CameraController@@UAEXXZ"));
    if (vtable == nullptr || get == nullptr || step == nullptr) {
        log::info("scene hooks: CameraController's vtable, getCameraTransform or step not found");
        return;
    }
    // The easing share: a float constant CameraController::step reads (at SceneWindow.dll's 0x100cb934 when loaded at
    // its preferred address), checked to be referenced by step's code and to be the expected 0.4.
    const auto* base = reinterpret_cast<const std::byte*>(sceneWindow);
    const auto* easing = reinterpret_cast<const float*>(base + 0xcb934);
    const auto address = static_cast<std::uint32_t>(reinterpret_cast<std::uintptr_t>(easing));
    bool referenced = false;
    for (std::size_t i = 0; i + 4 <= 0x700 && !referenced; ++i) {
        referenced = std::memcmp(step + i, &address, 4) == 0;
    }
    if (!referenced || *easing != 0.4f) {
        log::info("scene hooks: CameraController::step's easing not where expected");
        return;
    }
    cameraEasing = *easing;
    // CameraController's vtable: getCameraTransform is its 6th entry.
    originalGetCameraTransform = reinterpret_cast<GetTransform>(get);
    if (!replace(&vtable[5], reinterpret_cast<void*>(get), reinterpret_cast<void*>(&getCameraTransform))) {
        log::info("scene hooks: CameraController's vtable doesn't hold its getCameraTransform");
        return;
    }
    cameraHooked = 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) { // drawing between ticks needs the ticks
            install_pose(sceneWindow);
            install_particles(sceneWindow);
            install_camera(sceneWindow);
        }
        log::info("scene hooks: ticks {}, posing between ticks {}, particles between ticks {}, camera between ticks {}",
                  stepHooked ? "recorded" : "not recorded", poseHooked ? "on" : "off", particlesHooked ? "on" : "off",
                  cameraHooked ? "on" : "off");
    });
}

bool installed() { return stepHooked; }

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

} // namespace t5de::scene_hooks