#include "90CN_selected_page_fix.h"

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

namespace
{
// The cold-login failure occurs at the first field access inside
// sub_17D6BE0 when it is entered with ECX/this == NULL. Guard that function
// boundary itself instead of aborting the surrounding class1/op4 dispatcher;
// the latter owns additional HUD/state initialization that must still run.
constexpr uintptr_t kSelectedPageApplyRva = 0x013D6BE0; // 017D6BE0
constexpr size_t kHookLength = 5;
const unsigned char kSelectedPageApplyPrologue[kHookLength] = {
    0x55, 0x8B, 0xEC, 0x6A, 0xFF
};

HMODULE g_selfModule = nullptr;
uintptr_t g_exeBase = 0;
void* g_originalSelectedPageApply = nullptr;
volatile LONG g_bypassSequence = 0;

typedef void (__thiscall* SelectedPageApplyFn)(void*);

void BuildLogPath(wchar_t* output, size_t outputCount)
{
    if (!output || outputCount == 0) return;
    output[0] = L'\0';
    DWORD length = g_selfModule
        ? GetModuleFileNameW(g_selfModule, output, static_cast<DWORD>(outputCount))
        : 0;
    if (!length || length >= outputCount) {
        wcscpy_s(output, outputCount, L"90CN_fix.log");
        return;
    }

    wchar_t* slash = wcsrchr(output, L'\\');
    wchar_t* slash2 = wcsrchr(output, L'/');
    if (slash2 && (!slash || slash2 > slash)) slash = slash2;
    if (slash) {
        slash[1] = L'\0';
        wcscat_s(output, outputCount, L"90CN_fix.log");
    }
    else {
        wcscpy_s(output, outputCount, L"90CN_fix.log");
    }
}

void WriteLogLine(const char* text)
{
    if (!text) return;
    wchar_t path[MAX_PATH] = {};
    BuildLogPath(path, _countof(path));
    HANDLE file = CreateFileW(path, FILE_APPEND_DATA,
        FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
        nullptr, OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL, nullptr);
    if (file == INVALID_HANDLE_VALUE) return;
    DWORD written = 0;
    WriteFile(file, text, static_cast<DWORD>(strlen(text)), &written, nullptr);
    CloseHandle(file);
}

void RecordNullBypass()
{
    LONG sequence = InterlockedIncrement(&g_bypassSequence);
    if (sequence > 64) return;

    SYSTEMTIME now = {};
    GetLocalTime(&now);
    char line[512] = {};
    sprintf_s(line,
        "selected-page-readiness-fix seq=%ld time=%04u-%02u-%02u %02u:%02u:%02u.%03u "
        "tid=%lu action=skip-null-this target=017D6BE0 boundary=selected-page-entry\r\n",
        sequence,
        now.wYear, now.wMonth, now.wDay,
        now.wHour, now.wMinute, now.wSecond, now.wMilliseconds,
        GetCurrentThreadId());
    WriteLogLine(line);
}

void __fastcall ProxySelectedPageApply(void* self, void* /*unused*/)
{
    if (!self) {
        RecordNullBypass();
        return;
    }

    SelectedPageApplyFn original = reinterpret_cast<SelectedPageApplyFn>(
        g_originalSelectedPageApply);
    if (original) original(self);
}

bool BuildTrampoline(unsigned char* target)
{
    const size_t trampolineLength = kHookLength + 5;
    unsigned char* trampoline = reinterpret_cast<unsigned char*>(VirtualAlloc(
        nullptr, trampolineLength, MEM_COMMIT | MEM_RESERVE,
        PAGE_EXECUTE_READWRITE));
    if (!trampoline) return false;

    memcpy(trampoline, target, kHookLength);
    trampoline[kHookLength] = 0xE9;
    *reinterpret_cast<LONG*>(trampoline + kHookLength + 1) =
        static_cast<LONG>(
            reinterpret_cast<uintptr_t>(target + kHookLength) -
            reinterpret_cast<uintptr_t>(trampoline + trampolineLength));
    FlushInstructionCache(GetCurrentProcess(), trampoline, trampolineLength);
    g_originalSelectedPageApply = trampoline;
    return true;
}

DWORD WINAPI InstallSelectedPageReadinessFix(LPVOID)
{
    HMODULE executable = GetModuleHandleW(nullptr);
    g_exeBase = reinterpret_cast<uintptr_t>(executable);
    if (!g_exeBase || !g_selfModule) return 0;

    unsigned char* target = reinterpret_cast<unsigned char*>(
        g_exeBase + kSelectedPageApplyRva);

    bool verified = false;
    unsigned char observed[kHookLength] = {};
    __try {
        memcpy(observed, target, sizeof(observed));
        verified = memcmp(observed, kSelectedPageApplyPrologue,
            sizeof(kSelectedPageApplyPrologue)) == 0;
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        verified = false;
    }

    if (!verified) {
        char line[512] = {};
        sprintf_s(line,
            "selected-page-readiness-fix install_failed target=017D6BE0 "
            "observed=%02X%02X%02X%02X%02X expected=558BEC6AFF\r\n",
            observed[0], observed[1], observed[2], observed[3], observed[4]);
        WriteLogLine(line);
        return 0;
    }

    if (!BuildTrampoline(target)) {
        WriteLogLine(
            "selected-page-readiness-fix install_failed reason=trampoline-allocation\r\n");
        return 0;
    }

    const uintptr_t replacement = reinterpret_cast<uintptr_t>(
        &ProxySelectedPageApply);
    const LONG replacementRelative = static_cast<LONG>(
        replacement - reinterpret_cast<uintptr_t>(target + kHookLength));

    DWORD oldProtection = 0;
    if (!VirtualProtect(target, kHookLength, PAGE_EXECUTE_READWRITE,
            &oldProtection)) {
        VirtualFree(g_originalSelectedPageApply, 0, MEM_RELEASE);
        g_originalSelectedPageApply = nullptr;
        WriteLogLine(
            "selected-page-readiness-fix install_failed reason=VirtualProtect\r\n");
        return 0;
    }

    target[0] = 0xE9;
    *reinterpret_cast<LONG*>(target + 1) = replacementRelative;

    DWORD ignoredProtection = 0;
    VirtualProtect(target, kHookLength, oldProtection, &ignoredProtection);
    FlushInstructionCache(GetCurrentProcess(), target, kHookLength);

    WriteLogLine(
        "selected-page-readiness-fix installed boundary=selected-page-entry "
        "target=017D6BE0 scope=null-this-only preserves=class1-op4-continuation\r\n");
    return 0;
}
}

void StartSelectedPageReadinessFix(HMODULE module)
{
    g_selfModule = module;
    HANDLE worker = CreateThread(
        nullptr, 0, InstallSelectedPageReadinessFix, nullptr, 0, nullptr);
    if (worker) CloseHandle(worker);
}