// ==WindhawkMod==
// @id              classic-theme-windows
// @name            Classic Theme Windows
// @description     Forces Classic Theme on all Windows
// @version         0.4
// @author          Travis, Cynosphere
// @include         *
// @exclude         wininit.exe
// @exclude         winlogon.exe
// @exclude         taskmgr.exe
// @exclude         dwm.exe
// @exclude         C:\Windows\System32\*.scr
// @exclude         svchost.exe
// @exclude         taskhostw.exe
// @exclude         dllhost.exe
// @exclude         sihost.exe
// @exclude         lsass.exe
// @exclude         C:\Program Files (x86)\Steam\*
// @exclude         msedge.exe
// @exclude         vmware.exe
// @exclude         vmware-vmx.exe
// @exclude         Spotify.exe
// @exclude         smartscreen.exe
// @exclude         RuntimeBroker.exe
// @exclude         ApplicationFrameHost.exe
// @exclude         SystemSettings.exe
// @exclude         SecHealthUI.exe
// @exclude         SecurityHealthHost.exe
// @exclude         PhoneExperienceHost.exe
// @exclude         SecurityHealthTray.exe
// @exclude         Window Detective.exe
// @compilerOptions -luxtheme -ldwmapi -lole32 -lcomdlg32 -lcomctl32
// ==/WindhawkMod==

// ==WindhawkModReadme==
/*
# Classic Theme Windows
Forces Classic Theme on all Windows
*/
// ==/WindhawkModReadme==

#include <windhawk_api.h>
#include <windows.h>
#include <uxtheme.h>
#include <dwmapi.h>
#include <commctrl.h>

// credit to Anixx for this snippet
void ClientEdgeFixes(HWND hWnd) {
    HWND cn = NULL;
    HWND p = hWnd;
    wchar_t bufr[256];
    GetClassName(hWnd, bufr, 256);

    //Mathematica
    if (lstrcmpW(bufr, L"NotebookFrame") == 0)
    {
        p = FindWindowEx(p, NULL, L"NotebookContent", NULL);
        if(p != NULL)
        {
            cn = p;
        }

    }

    // Notepad
    if (lstrcmpW(bufr, L"Notepad") == 0)
    {
        p = FindWindowEx(p, NULL, L"Edit", NULL);
        if(p != NULL)
        {
            cn = p;
        }

    }

    // Console
    if (lstrcmpW(bufr, L"ConsoleWindowClass") == 0)
    {
        cn = p;
    }

    // File picker
    if (lstrcmpW(bufr, L"#32770") == 0)
    {
        p = FindWindowEx(p, NULL, L"SHELLDLL_DefView", NULL);
        if(p != NULL)
        {
            p = FindWindowEx(p, NULL, L"SysListView32", NULL);
            if(p != NULL)
            {
                cn = p;
            }
        }
    }
    
    if(cn != NULL)
    {
        SetWindowLongPtrW(cn, GWL_EXSTYLE, GetWindowLongPtrW(cn, GWL_EXSTYLE) | WS_EX_CLIENTEDGE);
        RECT rect;
        GetWindowRect(cn, &rect);
        SetWindowPos(cn, NULL, NULL, NULL, rect.right - rect.left + 1, rect.bottom - rect.top, SWP_NOMOVE | SWP_NOZORDER | SWP_DEFERERASE);
        SetWindowPos(cn, NULL, NULL, NULL, rect.right - rect.left, rect.bottom - rect.top, SWP_NOMOVE | SWP_NOZORDER | SWP_DEFERERASE);
    }
}

// BasicThemer2 reimplementation to render classic titlebars
static const int DisableDWM = DWMNCRP_DISABLED;
static const int EnableDWM = DWMNCRP_ENABLED;

void DisableDWMForHwnd(HWND hwnd) {
    DwmSetWindowAttribute(hwnd, DWMWA_NCRENDERING_POLICY, &DisableDWM, sizeof(int));
}
void EnableDWMForHwnd(HWND hwnd) {
    DwmSetWindowAttribute(hwnd, DWMWA_NCRENDERING_POLICY, &EnableDWM, sizeof(int));
}

DWORD g_uiThreadId;

BOOL CALLBACK DisableDWMForAll(HWND hWnd, LPARAM lParam) {
    DWORD dwProcessId = 0;
    DWORD dwThreadId = GetWindowThreadProcessId(hWnd, &dwProcessId);
    if (!dwThreadId || dwProcessId != GetCurrentProcessId()) {
        return TRUE;
    }

    if (g_uiThreadId && g_uiThreadId != dwThreadId) {
        return TRUE;
    }

    ClientEdgeFixes(hWnd);

    if (IsWindow(hWnd)) {
        //Wh_Log(L"[Init] Window found: %08X", (DWORD)(ULONG_PTR)hWnd);

        if (!g_uiThreadId) {
            g_uiThreadId = dwThreadId;
        }

        DisableDWMForHwnd(hWnd);
    }

    return TRUE;
}

BOOL CALLBACK EnableDWMForAll(HWND hWnd, LPARAM lParam) {
    DWORD dwProcessId = 0;
    DWORD dwThreadId = GetWindowThreadProcessId(hWnd, &dwProcessId);
    if (!dwThreadId || dwProcessId != GetCurrentProcessId()) {
        return TRUE;
    }

    if (g_uiThreadId && g_uiThreadId != dwThreadId) {
        return TRUE;
    }

    if (IsWindow(hWnd)) {
        //Wh_Log(L"[Cleanup] Window found: %08X", (DWORD)(ULONG_PTR)hWnd);

        if (!g_uiThreadId) {
            g_uiThreadId = dwThreadId;
        }

        EnableDWMForHwnd(hWnd);
    }

    return TRUE;
}

using CreateWindowExW_t = decltype(&CreateWindowExW);
CreateWindowExW_t CreateWindowExW_Orig;
HWND WINAPI CreateWindowExW_Hook(DWORD dwExStyle,LPCWSTR lpClassName,LPCWSTR lpWindowName,DWORD dwStyle,int X,int Y,int nWidth,int nHeight,HWND hWndParent,HMENU hMenu,HINSTANCE hInstance,LPVOID lpParam) {
    HWND hWnd = CreateWindowExW_Orig(dwExStyle,lpClassName,lpWindowName,dwStyle,X,Y,nWidth,nHeight,hWndParent,hMenu,hInstance,lpParam);

    if (!hWnd) {
        return hWnd;
    }

    if (g_uiThreadId && g_uiThreadId != GetCurrentThreadId()) {
        return hWnd;
    }

    if (IsWindow(hWnd)) {
        //Wh_Log(L"[CreateWindowExW] Window created: %08X", (DWORD)(ULONG_PTR)hWnd);

        if (!g_uiThreadId) {
            g_uiThreadId = GetCurrentThreadId();
        }

        DisableDWMForHwnd(hWnd);
        EnumWindows(DisableDWMForAll, 0);
    }

    return hWnd;
}
using CreateWindowExA_t = decltype(&CreateWindowExA);
CreateWindowExA_t CreateWindowExA_Orig;
HWND WINAPI CreateWindowExA_Hook(DWORD dwExStyle,LPCSTR lpClassName,LPCSTR lpWindowName,DWORD dwStyle,int X,int Y,int nWidth,int nHeight,HWND hWndParent,HMENU hMenu,HINSTANCE hInstance,LPVOID lpParam) {
    HWND hWnd = CreateWindowExA_Orig(dwExStyle,lpClassName,lpWindowName,dwStyle,X,Y,nWidth,nHeight,hWndParent,hMenu,hInstance,lpParam);

    if (!hWnd) {
        return hWnd;
    }

    if (g_uiThreadId && g_uiThreadId != GetCurrentThreadId()) {
        return hWnd;
    }

    if (IsWindow(hWnd)) {
        //Wh_Log(L"[CreateWindowExA] Window created: %08X", (DWORD)(ULONG_PTR)hWnd);

        if (!g_uiThreadId) {
            g_uiThreadId = GetCurrentThreadId();
        }

        DisableDWMForHwnd(hWnd);
        EnumWindows(DisableDWMForAll, 0);
    }

    return hWnd;
}

HWND(WINAPI *pOriginalSHCreateWorkerWindow)(WNDPROC wndProc, HWND hWndParent, DWORD dwExStyle, DWORD dwStyle, HMENU hMenu, LONG_PTR wnd_extra);
HWND WINAPI SHCreateWorkerWindowHook(WNDPROC wndProc, HWND hWndParent, DWORD dwExStyle, DWORD dwStyle, HMENU hMenu, LONG_PTR wnd_extra)
{
    HWND result = pOriginalSHCreateWorkerWindow(wndProc, hWndParent, dwExStyle, dwStyle, hMenu, wnd_extra);

    if (dwExStyle == 0x10000 && dwStyle == 0x46000000 && result) {
        DisableDWMForHwnd(hWndParent);
        EnumWindows(DisableDWMForAll, 0);
    }

    return result;
}

using ShowWindow_t = decltype(&ShowWindow);
ShowWindow_t pOriginalShowWindow;
WINBOOL WINAPI ShowWindowHook(HWND hWnd,int nCmdShow)
{
    bool ret = pOriginalShowWindow(hWnd, nCmdShow);

    if (g_uiThreadId && g_uiThreadId != GetCurrentThreadId()) {
        return ret;
    }

    if (IsWindow(hWnd)) {
        //Wh_Log(L"[ShowWindowHook] Window created: %08X", (DWORD)(ULONG_PTR)hWnd);

        if (!g_uiThreadId) {
            g_uiThreadId = GetCurrentThreadId();
        }

        DisableDWMForHwnd(hWnd);
        //EnumWindows(DisableDWMForAll, 0);
    }

    return ret;
}

// File picker
template<class S>
CLSID CreateGUID(const S& hexString)
{
    CLSID clsid;
    CLSIDFromString(hexString, &clsid);

    return clsid;
}

const CLSID CLSID_FileOpenDialog = CreateGUID(L"{DC1C5A9C-E88A-4dde-A5A1-60F82A20AEF7}");
const CLSID CLSID_FileOpenDialogLegacy = CreateGUID(L"{725F645B-EAED-4fc5-B1C5-D9AD0ACCBA5E}");
const CLSID CLSID_FileSaveDialog = CreateGUID(L"{c0b4e2f3-ba21-4773-8dba-335ec946eb8b}");
const CLSID CLSID_FileSaveDialogLegacy = CreateGUID(L"{AF02484C-A0A9-4669-9051-058AB12B9195}");

using CoCreateInstance_t = decltype(&CoCreateInstance);
CoCreateInstance_t CoCreateInstance_Original;
HRESULT WINAPI CoCreateInstance_Hook(REFCLSID rclsid, LPUNKNOWN pUnkOuter, DWORD dwClsContext, REFIID riid, LPVOID *ppv) {
    if (rclsid == CLSID_FileOpenDialog) {
        return CoCreateInstance_Original(CLSID_FileOpenDialogLegacy, pUnkOuter, dwClsContext, riid, ppv);
    } else if (rclsid == CLSID_FileSaveDialog) {
        return CoCreateInstance_Original(CLSID_FileSaveDialogLegacy, pUnkOuter, dwClsContext, riid, ppv);
    }

    return CoCreateInstance_Original(rclsid, pUnkOuter, dwClsContext, riid, ppv);
}

// Windhawk functions
BOOL Wh_ModInit() {
    // Enable classic theme
    SetThemeAppProperties(0);

    // BasicThemer hooks
    Wh_SetFunctionHook((void*)CreateWindowExW,
                       (void*)CreateWindowExW_Hook,
                       (void**)&CreateWindowExW_Orig);
    Wh_SetFunctionHook((void*)CreateWindowExA,
                       (void*)CreateWindowExA_Hook,
                       (void**)&CreateWindowExA_Orig);
    
    HMODULE hShcore = GetModuleHandle(L"shcore.dll");
    void* origFunc = (void*)GetProcAddress(hShcore, (LPCSTR)188);
    Wh_SetFunctionHook(origFunc, (void*)SHCreateWorkerWindowHook, (void**)&pOriginalSHCreateWorkerWindow);

    Wh_SetFunctionHook((void*)ShowWindow, (void*)ShowWindowHook, (void**)&pOriginalShowWindow);

    // Iterate every window to enable BasicThemer
    EnumWindows(DisableDWMForAll, 0);

    // File picker
    Wh_SetFunctionHook((void*)CoCreateInstance, (void*)CoCreateInstance_Hook,
                       (void**)&CoCreateInstance_Original);

    return TRUE;
}

void Wh_ModUninit() {
    // Disable classic theme
    if (GetThemeAppProperties() == 0) {
        SetThemeAppProperties(3);
    }

    // Iterate every window to disable BasicThemer
    EnumWindows(EnableDWMForAll, 0);
}