#include "90CN_3149_writer_resync_fix.h"

#include <stdint.h>
#include <string.h>

namespace
{
// Runtime-verified BAD path:
//   01B947C0 self-removal branch
//   01B947D6 push 0
//   01B947D8 call [vtable+6AC]
//   return 01B947DA
// On the affected retained transition, that virtual call resolves to
// 01B9A1A0(0), which directly clears both tier1 +3D1 bytes. The last tier1
// +0x28 update has already happened by then, so manager mode remains 0.
constexpr uintptr_t kSelfRemovalWriterCallsiteRva = 0x017947D6; // VA 01B947D6
constexpr uintptr_t kSelfRemovalWriterContinueRva = 0x017947DC; // VA 01B947DC

// 01B9A1A0 direct +3D1 writes for slot1 and slot0.
constexpr uintptr_t kWriterBSlot1WriteRva = 0x0179A275; // VA 01B9A275
constexpr uintptr_t kWriterBSlot0WriteRva = 0x0179A35D; // VA 01B9A35D
constexpr uintptr_t kWriterBSlot1ContinueRva = 0x0179A27B;
constexpr uintptr_t kWriterBSlot0ContinueRva = 0x0179A363;

constexpr uintptr_t kSceneRootPointerRva = 0x04DB2764;
constexpr uintptr_t kSceneUiManagerRva = 0x01490EE0;
constexpr size_t kD1Offset = 0x3D1;
constexpr size_t kModeOffset = 0x719;

uintptr_t g_exeBase = 0;
uintptr_t g_callsiteContinue = 0;
uintptr_t g_slot1Continue = 0;
uintptr_t g_slot0Continue = 0;

volatile LONG g_captureActive = 0;
volatile LONG g_captureThreadId = 0;
uintptr_t g_captureFirst = 0;
uintptr_t g_captureSecond = 0;

typedef void* (__thiscall* SceneUiManagerFn)(void*);
typedef void (__thiscall* Tier1UpdateFn)(void*);

unsigned int ReadManagerMode(bool* fault)
{
    if (fault) *fault = false;
    unsigned int mode = 0xFFFFFFFFu;
    __try {
        void* scene = *reinterpret_cast<void**>(g_exeBase + kSceneRootPointerRva);
        if (!scene) return mode;
        SceneUiManagerFn getManager = reinterpret_cast<SceneUiManagerFn>(
            g_exeBase + kSceneUiManagerRva);
        void* manager = getManager ? getManager(scene) : nullptr;
        if (!manager) return mode;
        mode = *reinterpret_cast<unsigned char*>(
            reinterpret_cast<uintptr_t>(manager) + kModeOffset);
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        if (fault) *fault = true;
        mode = 0xFFFFFFFFu;
    }
    return mode;
}

unsigned int ReadD1(uintptr_t object, bool* fault)
{
    if (fault) *fault = false;
    if (!object) return 0xFFFFFFFFu;
    unsigned int value = 0xFFFFFFFFu;
    __try {
        value = *reinterpret_cast<unsigned char*>(object + kD1Offset);
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        if (fault) *fault = true;
    }
    return value;
}

bool PatchJumpExact(uintptr_t target, const unsigned char* expected, size_t patchLength,
    void* detour, uintptr_t continuation)
{
    if (!target || !expected || patchLength < 5 || !detour || !continuation) return false;
    __try {
        if (memcmp(reinterpret_cast<const void*>(target), expected, patchLength) != 0) {
            return false;
        }
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        return false;
    }

    DWORD oldProtection = 0;
    if (!VirtualProtect(reinterpret_cast<void*>(target), patchLength,
            PAGE_EXECUTE_READWRITE, &oldProtection)) {
        return false;
    }

    unsigned char* site = reinterpret_cast<unsigned char*>(target);
    site[0] = 0xE9;
    *reinterpret_cast<LONG*>(site + 1) = static_cast<LONG>(
        reinterpret_cast<uintptr_t>(detour) - (target + 5));
    for (size_t i = 5; i < patchLength; ++i) site[i] = 0x90;

    DWORD ignoredProtection = 0;
    VirtualProtect(site, patchLength, oldProtection, &ignoredProtection);
    FlushInstructionCache(GetCurrentProcess(), site, patchLength);
    return true;
}

void __stdcall BeginTier1Capture()
{
    g_captureFirst = 0;
    g_captureSecond = 0;
    InterlockedExchange(&g_captureThreadId, static_cast<LONG>(GetCurrentThreadId()));
    InterlockedExchange(&g_captureActive, 1);
}

void __stdcall CaptureTier1Object(void* object)
{
    if (!object) return;
    if (InterlockedCompareExchange(&g_captureActive, 1, 1) != 1) return;
    if (static_cast<DWORD>(InterlockedCompareExchange(&g_captureThreadId, 0, 0)) !=
        GetCurrentThreadId()) {
        return;
    }

    uintptr_t address = reinterpret_cast<uintptr_t>(object);
    if (!g_captureFirst) {
        g_captureFirst = address;
    }
    else if (g_captureFirst != address && !g_captureSecond) {
        g_captureSecond = address;
    }
}

bool CallTier1Virtual28(uintptr_t object)
{
    if (!object) return false;
    bool ok = false;
    __try {
        uintptr_t* vtable = *reinterpret_cast<uintptr_t**>(object);
        if (!vtable) return false;
        uintptr_t target = vtable[0x28 / sizeof(uintptr_t)];
        if (!target) return false;
        Tier1UpdateFn update = reinterpret_cast<Tier1UpdateFn>(target);
        update(reinterpret_cast<void*>(object));
        ok = true;
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        ok = false;
    }
    return ok;
}

void __stdcall FinishTier1CaptureAndResync()
{
    if (static_cast<DWORD>(InterlockedCompareExchange(&g_captureThreadId, 0, 0)) !=
        GetCurrentThreadId()) {
        return;
    }
    if (InterlockedExchange(&g_captureActive, 0) != 1) return;

    uintptr_t first = g_captureFirst;
    uintptr_t second = g_captureSecond;
    g_captureFirst = 0;
    g_captureSecond = 0;

    bool modeFault = false;
    unsigned int mode = ReadManagerMode(&modeFault);
    bool d1FaultFirst = false;
    bool d1FaultSecond = false;
    unsigned int firstD1 = ReadD1(first, &d1FaultFirst);
    unsigned int secondD1 = ReadD1(second, &d1FaultSecond);

    // Scope the behavior change to the exact stale-HUD state observed live:
    // manager mode is 0 after the writer has cleared at least one captured
    // tier1 object to +3D1 == 0. If mode is already 1, do nothing.
    if (!modeFault && mode == 0) {
        if (first && !d1FaultFirst && firstD1 == 0) {
            (void)CallTier1Virtual28(first);
        }
        if (second && second != first && !d1FaultSecond && secondD1 == 0) {
            (void)CallTier1Virtual28(second);
        }
    }
}

__declspec(naked) void ProxyWriterBSlot1Write()
{
    __asm {
        mov byte ptr [esi + 3D1h], bl
        pushfd
        pushad
        push esi
        call CaptureTier1Object
        popad
        popfd
        jmp dword ptr [g_slot1Continue]
    }
}

__declspec(naked) void ProxyWriterBSlot0Write()
{
    __asm {
        mov byte ptr [esi + 3D1h], bl
        pushfd
        pushad
        push esi
        call CaptureTier1Object
        popad
        popfd
        jmp dword ptr [g_slot0Continue]
    }
}

__declspec(naked) void ProxySelfRemovalWriterZero()
{
    __asm {
        // Reproduce 01B947D6..01B947DA exactly around the native virtual call:
        // push 0; call edx; [post-call resync]; push 1; continue at 01B947DC.
        pushfd
        pushad
        call BeginTier1Capture
        popad
        popfd

        push 0
        call edx

        pushfd
        pushad
        call FinishTier1CaptureAndResync
        popad
        popfd

        push 1
        jmp dword ptr [g_callsiteContinue]
    }
}

DWORD WINAPI Install3149WriterResyncFix(LPVOID)
{
    g_exeBase = reinterpret_cast<uintptr_t>(GetModuleHandleW(nullptr));
    if (!g_exeBase) return 0;

    g_callsiteContinue = g_exeBase + kSelfRemovalWriterContinueRva;
    g_slot1Continue = g_exeBase + kWriterBSlot1ContinueRva;
    g_slot0Continue = g_exeBase + kWriterBSlot0ContinueRva;

    static const unsigned char kWriterWriteExpected[] = {
        0x88, 0x9E, 0xD1, 0x03, 0x00, 0x00
    };
    static const unsigned char kSelfRemovalExpected[] = {
        0x6A, 0x00, 0xFF, 0xD2, 0x6A, 0x01
    };

    bool slot1Hook = PatchJumpExact(
        g_exeBase + kWriterBSlot1WriteRva,
        kWriterWriteExpected, sizeof(kWriterWriteExpected),
        reinterpret_cast<void*>(&ProxyWriterBSlot1Write), g_slot1Continue);
    bool slot0Hook = PatchJumpExact(
        g_exeBase + kWriterBSlot0WriteRva,
        kWriterWriteExpected, sizeof(kWriterWriteExpected),
        reinterpret_cast<void*>(&ProxyWriterBSlot0Write), g_slot0Continue);

    // Only enable the behavior-changing callsite hook if both read/capture
    // hooks are present. Partial capture must never drive a resync.
    if (slot1Hook && slot0Hook) {
        (void)PatchJumpExact(
            g_exeBase + kSelfRemovalWriterCallsiteRva,
            kSelfRemovalExpected, sizeof(kSelfRemovalExpected),
            reinterpret_cast<void*>(&ProxySelfRemovalWriterZero), g_callsiteContinue);
    }
    return 0;
}
}

void Start3149WriterResyncFix(HMODULE module)
{
    (void)module;
    HANDLE worker = CreateThread(nullptr, 0, Install3149WriterResyncFix, nullptr, 0, nullptr);
    if (worker) CloseHandle(worker);
}