/***
    Copyright (C) 2021 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: DarkTheme.cpp
    Date: 2021-6-11
    Author: Reece
***/
#include <RuntimeInternal.hpp>
#include "DarkTheme.hpp"
#include <Uxtheme.h>

namespace Aurora::Extensions::Win32
{
    enum IMMERSIVE_HC_CACHE_MODE
    {
        IHCM_USE_CACHED_VALUE,
        IHCM_REFRESH
    };

    enum PreferredAppMode
    {
        Default,
        AllowDark,
        ForceDark,
        ForceLight,
        Max
    };

    enum WINDOWCOMPOSITIONATTRIB
    {
        WCA_UNDEFINED = 0,
        WCA_USEDARKMODECOLORS = 26,
        WCA_LAST = 27
    };

    struct WINDOWCOMPOSITIONATTRIBDATA
    {
        WINDOWCOMPOSITIONATTRIB Attrib;
        PVOID pvData;
        SIZE_T cbData;
    };

    using fnRtlGetNtVersionNumbers = void (WINAPI *)(LPDWORD major, LPDWORD minor, LPDWORD build);
    using fnSetWindowCompositionAttribute = BOOL(WINAPI *)(HWND hWnd, WINDOWCOMPOSITIONATTRIBDATA *);

    // 1809 17763
    using fnShouldAppsUseDarkMode   = bool (WINAPI *)(); // ordinal 132
    using fnAllowDarkModeForWindow  = bool (WINAPI *)(HWND hWnd, bool allow); // ordinal 133
    using fnAllowDarkModeForApp     = bool (WINAPI *)(bool allow); // ordinal 135, in 1809

    using fnSetPreferredAppMode     = PreferredAppMode(WINAPI *)(PreferredAppMode appMode); // ordinal 135, in 1903
    using fnIsDarkModeAllowedForApp = bool (WINAPI *)(); // ordinal 139

    static fnSetWindowCompositionAttribute _SetWindowCompositionAttribute = nullptr;
    static fnAllowDarkModeForWindow        _AllowDarkModeForWindow = nullptr;
    static fnAllowDarkModeForApp           _AllowDarkModeForApp = nullptr;
    // 1903 18362                    
    static fnSetPreferredAppMode           _SetPreferredAppMode = nullptr;

    static DWORD g_buildNumber = 0;

    bool AllowDarkModeForWindow(HWND hWnd, bool allow)
    {
        if (g_darkModeSupported)
        {
            return _AllowDarkModeForWindow(hWnd, allow);
        }
        return false;
    }

    void RefreshTitleBarThemeColor(HWND hWnd)
    {
        BOOL dark = TRUE;
        
        if (g_buildNumber < 18362)
        {
            SetPropW(hWnd, L"UseImmersiveDarkModeColors", reinterpret_cast<HANDLE>(static_cast<INT_PTR>(dark)));
        }
        else if (_SetWindowCompositionAttribute)
        {
            WINDOWCOMPOSITIONATTRIBDATA data = { WCA_USEDARKMODECOLORS, &dark, sizeof(dark) };
            _SetWindowCompositionAttribute(hWnd, &data);
        }
    }

    void MakeWindowBordersDark(HWND window)
    {
        if (!g_darkModeSupported) return;
        SetWindowTheme(window, L"DarkMode_Explorer", NULL);
        AllowDarkModeForWindow(window, true);
        RefreshTitleBarThemeColor(window);
    }

    void AllowDarkModeForApp()
    {
        if (_AllowDarkModeForApp)
        {
            _AllowDarkModeForApp(true);
        }
        else if (_SetPreferredAppMode)
        {
            _SetPreferredAppMode(ForceDark);
        }
    }

    static constexpr bool CheckBuildNumber(DWORD buildNumber)
    {
        return (buildNumber == 17763 || // 1809
                buildNumber == 18362 || // 1903
                buildNumber == 18363 || // 1909
                buildNumber == 19041);  // 2004
    }

    using fnRefreshImmersiveColorPolicyState = void (WINAPI *)(); // ordinal 104
    fnRefreshImmersiveColorPolicyState _RefreshImmersiveColorPolicyState = nullptr;

    void InitDarkMode()
    {
        auto RtlGetNtVersionNumbers = reinterpret_cast<fnRtlGetNtVersionNumbers>(GetProcAddress(GetModuleHandleW(L"ntdll.dll"), "RtlGetNtVersionNumbers"));
        if (!RtlGetNtVersionNumbers)
        {
            return;
        }
        
        DWORD major, minor;
        RtlGetNtVersionNumbers(&major, &minor, &g_buildNumber);
        g_buildNumber &= ~0xF0000000;
        if (major != 10 || minor != 0 || !CheckBuildNumber(g_buildNumber))
        {
            return;
        }

        HMODULE hUxtheme = LoadLibraryExW(L"uxtheme.dll", nullptr, LOAD_LIBRARY_SEARCH_SYSTEM32);
        if (!hUxtheme)
        {
            return;
        }

        _AllowDarkModeForWindow = reinterpret_cast<fnAllowDarkModeForWindow>(GetProcAddress(hUxtheme, MAKEINTRESOURCEA(133)));
        _RefreshImmersiveColorPolicyState = reinterpret_cast<fnRefreshImmersiveColorPolicyState>(GetProcAddress(hUxtheme, MAKEINTRESOURCEA(104)));
        auto ord135 = GetProcAddress(hUxtheme, MAKEINTRESOURCEA(135));

        if (g_buildNumber < 18362)
        {
            _AllowDarkModeForApp = reinterpret_cast<fnAllowDarkModeForApp>(ord135);
        }
        else
        {
            _SetPreferredAppMode = reinterpret_cast<fnSetPreferredAppMode>(ord135);
        }

        _FlushMenuThemes = reinterpret_cast<fnFlushMenuThemes>(GetProcAddress(hUxtheme, MAKEINTRESOURCEA(136)));

        _SetWindowCompositionAttribute = reinterpret_cast<fnSetWindowCompositionAttribute>(GetProcAddress(GetModuleHandleW(L"user32.dll"), "SetWindowCompositionAttribute"));

        if (_AllowDarkModeForWindow &&
            (_AllowDarkModeForApp || _SetPreferredAppMode))
        {
            g_darkModeSupported = true;

            AllowDarkModeForApp();
            if (_RefreshImmersiveColorPolicyState)
            {
                if (g_buildNumber < 18362)
                {
                    _RefreshImmersiveColorPolicyState();
                }
            }
        }
    }
}