#include "90CN_weapon_skill_relogin_fix.h"

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

namespace
{
constexpr uintptr_t kRequirementBranchRva = 0x01C07739;
constexpr uintptr_t kRequirementInvalidResumeRva = 0x01C07748;
constexpr uintptr_t kRequirementValidResumeRva = 0x01C07780;
constexpr uintptr_t kNativeSkillValidityRva = 0x0225F610;
constexpr uintptr_t kReloginEquipmentCallerReturnRva = 0x01C0859E;
constexpr DWORD kInstallDelayMs = 1500;

// Native bytes immediately after actor->lookup(requirementId, 1):
//   je invalid
//   push 1
//   mov ecx,eax
const unsigned char kExpectedBytes[6] = { 0x74, 0x0D, 0x6A, 0x01, 0x8B, 0xC8 };

HMODULE g_selfModule = nullptr;
uintptr_t g_exeBase = 0;
HANDLE g_worker = nullptr;
volatile LONG g_stop = 0;
volatile LONG g_installed = 0;
volatile LONG g_bypassLogCount = 0;
SRWLOCK g_logLock = SRWLOCK_INIT;
void* g_nativeSkillValidity = nullptr;
void* g_invalidResume = nullptr;
void* g_validResume = nullptr;
uintptr_t g_expectedCallerReturn = 0;

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 WriteLog(const char* format, ...)
{
    if (!format) return;
    char message[2048] = {};
    va_list args;
    va_start(args, format);
    _vsnprintf_s(message, sizeof(message), _TRUNCATE, format, args);
    va_end(args);

    SYSTEMTIME now = {};
    GetLocalTime(&now);
    char line[2304] = {};
    _snprintf_s(line, sizeof(line), _TRUNCATE,
        "%04u-%02u-%02u %02u:%02u:%02u.%03u tid=%lu %s\r\n",
        now.wYear, now.wMonth, now.wDay,
        now.wHour, now.wMinute, now.wSecond, now.wMilliseconds,
        static_cast<unsigned long>(GetCurrentThreadId()), message);

    wchar_t path[MAX_PATH] = {};
    BuildLogPath(path, _countof(path));
    AcquireSRWLockExclusive(&g_logLock);
    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) {
        DWORD written = 0;
        WriteFile(file, line, static_cast<DWORD>(strlen(line)), &written, nullptr);
        CloseHandle(file);
    }
    ReleaseSRWLockExclusive(&g_logLock);
}

__declspec(noinline) void __cdecl LogReloginBypass(
    unsigned int requirementId, uintptr_t callerReturn)
{
    LONG count = InterlockedIncrement(&g_bypassLogCount);
    if (count > 8) return;
    WriteLog(
        "weapon-relogin-fix bypass hit=%ld requirement_id=%u caller_rva=0x%08X "
        "lookup_record=present native_validity=0 action=treat-valid "
        "scope=op24-relogin-main-weapon-only",
        count, requirementId,
        g_exeBase && callerReturn >= g_exeBase
            ? static_cast<unsigned int>(callerReturn - g_exeBase)
            : 0u);
}

// Original sequence at RVA 0x01C07737:
//   test eax,eax
//   je   invalid
//   push 1
//   mov  ecx,eax
//   call sub_265F610
//   test eax,eax
//   jne  valid
//
// The caller has already decoded the equipped main weapon's requirement id,
// skipped requirementId == -1, and resolved actor->lookup(requirementId, 1).
// Live evidence proved that op24 can reach this point with a real requirement
// record while sub_265F610 transiently returns 0 before skill initialization is
// fully settled. Preserve every native decision except that timing window.
//
// The standard EBP frame lets us require [EBP+4] == base + 0x01C0859E, the
// observed op24 relogin main-weapon revalidation caller. Therefore normal
// manual equipping, other callers, missing requirement records, and native
// validity successes all keep their original behavior. No skill/profession
// allow list is needed, so the same mechanism covers other exclusive weapons.
__declspec(naked) void ProxyWeaponSkillReloginFix()
{
    __asm {
        // EAX is actor->lookup(requirementId, 1). A missing record remains
        // invalid exactly as in native code.
        test eax, eax
        jz requirement_invalid

        // Preserve the native validity call and accept native success first.
        push 1
        mov ecx, eax
        call dword ptr [g_nativeSkillValidity]
        test eax, eax
        jnz requirement_valid

        // Only repair the known relogin main-weapon revalidation call site.
        // EBX still holds the decoded requirement id for logging.
        mov edx, dword ptr [ebp + 4]
        cmp edx, dword ptr [g_expectedCallerReturn]
        jne requirement_invalid

        pushfd
        pushad
        push edx
        push ebx
        call LogReloginBypass
        add esp, 8
        popad
        popfd

        mov eax, 1
        jmp dword ptr [g_validResume]

    requirement_invalid:
        jmp dword ptr [g_invalidResume]

    requirement_valid:
        jmp dword ptr [g_validResume]
    }
}

bool PatchFix(unsigned char* target)
{
    if (!target) return false;
    unsigned char patch[6] = { 0xE9, 0, 0, 0, 0, 0x90 };
    LONG relative = static_cast<LONG>(
        reinterpret_cast<uintptr_t>(&ProxyWeaponSkillReloginFix) -
        reinterpret_cast<uintptr_t>(target + 5));
    *reinterpret_cast<LONG*>(patch + 1) = relative;

    DWORD oldProtection = 0;
    if (!VirtualProtect(target, sizeof(patch), PAGE_EXECUTE_READWRITE,
            &oldProtection)) {
        return false;
    }
    memcpy(target, patch, sizeof(patch));
    DWORD ignored = 0;
    VirtualProtect(target, sizeof(patch), oldProtection, &ignored);
    FlushInstructionCache(GetCurrentProcess(), target, sizeof(patch));
    return memcmp(target, patch, sizeof(patch)) == 0;
}

bool RestoreFix(unsigned char* target)
{
    if (!target) return false;
    DWORD oldProtection = 0;
    if (!VirtualProtect(target, sizeof(kExpectedBytes), PAGE_EXECUTE_READWRITE,
            &oldProtection)) {
        return false;
    }
    memcpy(target, kExpectedBytes, sizeof(kExpectedBytes));
    DWORD ignored = 0;
    VirtualProtect(target, sizeof(kExpectedBytes), oldProtection, &ignored);
    FlushInstructionCache(GetCurrentProcess(), target, sizeof(kExpectedBytes));
    return memcmp(target, kExpectedBytes, sizeof(kExpectedBytes)) == 0;
}

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

    DWORD slept = 0;
    while (slept < kInstallDelayMs &&
        InterlockedCompareExchange(&g_stop, 0, 0) == 0) {
        Sleep(100);
        slept += 100;
    }
    if (InterlockedCompareExchange(&g_stop, 0, 0) != 0) return 0;

    unsigned char* target = reinterpret_cast<unsigned char*>(
        g_exeBase + kRequirementBranchRva);
    g_nativeSkillValidity = reinterpret_cast<void*>(
        g_exeBase + kNativeSkillValidityRva);
    g_invalidResume = reinterpret_cast<void*>(
        g_exeBase + kRequirementInvalidResumeRva);
    g_validResume = reinterpret_cast<void*>(
        g_exeBase + kRequirementValidResumeRva);
    g_expectedCallerReturn = g_exeBase + kReloginEquipmentCallerReturnRva;

    __try {
        if (memcmp(target, kExpectedBytes, sizeof(kExpectedBytes)) != 0) {
            WriteLog(
                "weapon-relogin-fix install-failed target_rva=0x%08X "
                "observed=%02X%02X%02X%02X%02X%02X expected=740D6A018BC8",
                static_cast<unsigned int>(kRequirementBranchRva),
                target[0], target[1], target[2], target[3], target[4], target[5]);
            return 0;
        }
        if (!PatchFix(target)) {
            WriteLog(
                "weapon-relogin-fix install-failed reason=patch-write target_rva=0x%08X",
                static_cast<unsigned int>(kRequirementBranchRva));
            return 0;
        }
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        WriteLog("weapon-relogin-fix install exception code=0x%08X",
            GetExceptionCode());
        return 0;
    }

    InterlockedExchange(&g_installed, 1);
    WriteLog(
        "weapon-relogin-fix installed target_rva=0x%08X native_validity_rva=0x%08X "
        "caller_rva=0x%08X requirement_scope=any-present-record "
        "scope=op24-relogin-main-weapon-only",
        static_cast<unsigned int>(kRequirementBranchRva),
        static_cast<unsigned int>(kNativeSkillValidityRva),
        static_cast<unsigned int>(kReloginEquipmentCallerReturnRva));
    return 0;
}
}

void StartWeaponSkillReloginFix(HMODULE module)
{
    if (g_worker || !module) return;
    g_selfModule = module;
    InterlockedExchange(&g_stop, 0);
    g_worker = CreateThread(nullptr, 0, InstallWorker, nullptr, 0, nullptr);
    if (!g_worker) {
        WriteLog("weapon-relogin-fix start failed reason=CreateThread error=%lu",
            GetLastError());
    }
}

void StopWeaponSkillReloginFix()
{
    InterlockedExchange(&g_stop, 1);
    if (g_worker) {
        WaitForSingleObject(g_worker, 1500);
        CloseHandle(g_worker);
        g_worker = nullptr;
    }
    if (InterlockedCompareExchange(&g_installed, 0, 0) != 0 && g_exeBase) {
        unsigned char* target = reinterpret_cast<unsigned char*>(
            g_exeBase + kRequirementBranchRva);
        bool restored = false;
        __try {
            restored = RestoreFix(target);
        }
        __except (EXCEPTION_EXECUTE_HANDLER) {
            restored = false;
        }
        WriteLog("weapon-relogin-fix restore result=%d", restored ? 1 : 0);
        InterlockedExchange(&g_installed, 0);
    }
}