// ==WindhawkMod==
// @id              snippingtool-framerate
// @name            Snipping Tool Frame Rate Override
// @description     Overrides the minimum frame rate for Snipping Tool screen recording
// @version         1.1.0
// @author          Alchemy
// @include         SnippingTool.exe
// @architecture    x86-64
// ==/WindhawkMod==

// ==WindhawkModSettings==
/*
- frameIntervalMs: 8
  $name: Frame Interval (ms)
  $description: Milliseconds between frames. Lower = higher FPS. Default is 16ms (~60fps). Set to 8 for ~120fps, 4 for ~240fps.
*/
// ==/WindhawkModSettings==

#include <windhawk_utils.h>
#include <psapi.h>

struct {
    int frameIntervalMs;
} g_settings;

// Pattern for VideoFrameGenerator initialization function (sub_14006E680)
// Function prologue + unique offset references
// 48 89 5c 24 10       mov [rsp+10h], rbx
// 48 89 6c 24 18       mov [rsp+18h], rbp  
// 56                   push rsi
// 57                   push rdi
// 41 56                push r14
// 48 83 ec 60          sub rsp, 60h
// ...followed by unique pattern with offset 0x148 and 0x149 checks

// Signature pattern with wildcards (? = any byte)
// We use the prologue + distinctive offset pattern that checks member variables
const BYTE g_pattern[] = {
    0x48, 0x89, 0x5C, 0x24, 0x10,  // mov [rsp+10h], rbx
    0x48, 0x89, 0x6C, 0x24, 0x18,  // mov [rsp+18h], rbp
    0x56,                          // push rsi
    0x57,                          // push rdi
    0x41, 0x56,                    // push r14
    0x48, 0x83, 0xEC, 0x60,        // sub rsp, 60h
};

const char g_patternMask[] = "xxxxxxxxxxxxxxxxxx";

// Secondary pattern to verify we're in the right function
// This checks for the "cmp byte ptr [rbx+149h], 0" instruction
// 80 BB 49 01 00 00 00 = cmp byte ptr [rbx+149h], 0
const BYTE g_verifyPattern[] = {
    0x80, 0xBB, 0x49, 0x01, 0x00, 0x00, 0x00
};

using VideoFrameGenerator_Init_t = __int64 (__fastcall*)(
    __int64 thisPtr,
    char flag,
    __int64* mediaPropsFrameRate
);
VideoFrameGenerator_Init_t VideoFrameGenerator_Init_Original;

__int64 __fastcall VideoFrameGenerator_Init_Hook(
    __int64 thisPtr,
    char flag,
    __int64* mediaPropsFrameRate
) {
    // Call the original function first
    __int64 result = VideoFrameGenerator_Init_Original(thisPtr, flag, mediaPropsFrameRate);
    
    // Override m_minFrameRateMs at offset 432 (0x1B0)
    __int64* pFrameRateMs = (__int64*)(thisPtr + 432);
    
    int originalValue = (int)*pFrameRateMs;
    *pFrameRateMs = g_settings.frameIntervalMs;
    
    Wh_Log(L"VideoFrameGenerator: Overrode m_minFrameRateMs from %d to %d ms (~%d FPS)",
           originalValue,
           g_settings.frameIntervalMs,
           1000 / g_settings.frameIntervalMs);
    
    return result;
}

void* FindPattern(BYTE* base, SIZE_T size, const BYTE* pattern, const char* mask) {
    SIZE_T patternLen = strlen(mask);
    
    for (SIZE_T i = 0; i <= size - patternLen; i++) {
        bool found = true;
        for (SIZE_T j = 0; j < patternLen; j++) {
            if (mask[j] == 'x' && base[i + j] != pattern[j]) {
                found = false;
                break;
            }
        }
        if (found) {
            return base + i;
        }
    }
    return nullptr;
}

bool VerifyFunction(BYTE* funcAddr) {
    // Check if the verify pattern exists within the first 0x100 bytes
    for (int i = 0; i < 0x100 - sizeof(g_verifyPattern); i++) {
        bool match = true;
        for (size_t j = 0; j < sizeof(g_verifyPattern); j++) {
            if (funcAddr[i + j] != g_verifyPattern[j]) {
                match = false;
                break;
            }
        }
        if (match) {
            return true;
        }
    }
    return false;
}

void LoadSettings() {
    g_settings.frameIntervalMs = Wh_GetIntSetting(L"frameIntervalMs");
    
    if (g_settings.frameIntervalMs < 1) {
        g_settings.frameIntervalMs = 1;
    }
    if (g_settings.frameIntervalMs > 1000) {
        g_settings.frameIntervalMs = 1000;
    }
    
    Wh_Log(L"Settings loaded: frameIntervalMs = %d (~%d FPS)",
           g_settings.frameIntervalMs,
           1000 / g_settings.frameIntervalMs);
}

BOOL Wh_ModInit() {
    Wh_Log(L"Snipping Tool Frame Rate Override initializing");
    
    LoadSettings();
    
    HMODULE hModule = GetModuleHandleW(nullptr);
    if (!hModule) {
        Wh_Log(L"Failed to get module handle");
        return FALSE;
    }
    
    // Get module info for pattern scanning
    MODULEINFO modInfo;
    if (!GetModuleInformation(GetCurrentProcess(), hModule, &modInfo, sizeof(modInfo))) {
        Wh_Log(L"Failed to get module information");
        return FALSE;
    }
    
    BYTE* moduleBase = (BYTE*)modInfo.lpBaseOfDll;
    SIZE_T moduleSize = modInfo.SizeOfImage;
    
    Wh_Log(L"Scanning module: base=%p, size=0x%X", moduleBase, (DWORD)moduleSize);
    
    // Scan for the pattern
    void* foundAddr = nullptr;
    BYTE* searchStart = moduleBase;
    SIZE_T searchSize = moduleSize;
    
    while (searchSize > sizeof(g_pattern)) {
        void* candidate = FindPattern(searchStart, searchSize, g_pattern, g_patternMask);
        
        if (!candidate) {
            break;
        }
        
        Wh_Log(L"Found pattern candidate at %p", candidate);
        
        // Verify this is the correct function
        if (VerifyFunction((BYTE*)candidate)) {
            foundAddr = candidate;
            Wh_Log(L"Verified! This is the VideoFrameGenerator init function");
            break;
        }
        
        // Continue searching after this candidate
        BYTE* nextSearch = (BYTE*)candidate + 1;
        searchSize = moduleSize - (nextSearch - moduleBase);
        searchStart = nextSearch;
    }
    
    if (!foundAddr) {
        Wh_Log(L"Pattern not found, trying fallback offset 0x6E680");
        
        // Fallback to known offset
        foundAddr = moduleBase + 0x6E680;
        
        // Verify the fallback
        bool validFallback = true;
        for (size_t i = 0; i < sizeof(g_pattern); i++) {
            if (((BYTE*)foundAddr)[i] != g_pattern[i]) {
                validFallback = false;
                break;
            }
        }
        
        if (!validFallback) {
            Wh_Log(L"Fallback offset also invalid. This version may not be supported.");
            return FALSE;
        }
        
        Wh_Log(L"Fallback offset verified at %p", foundAddr);
    }
    
    // Set the hook
    if (!Wh_SetFunctionHook(
            foundAddr,
            (void*)VideoFrameGenerator_Init_Hook,
            (void**)&VideoFrameGenerator_Init_Original)) {
        Wh_Log(L"Failed to set function hook");
        return FALSE;
    }
    
    Wh_Log(L"Successfully hooked VideoFrameGenerator::Initialize at %p", foundAddr);
    return TRUE;
}

void Wh_ModSettingsChanged() {
    Wh_Log(L"Settings changed");
    LoadSettings();
}

void Wh_ModUninit() {
    Wh_Log(L"Snipping Tool Frame Rate Override unloading");
}
