#include "pch.h"
#include "Helpers.h"
#include <sstream>

#include <common/interop/shared_constants.h>
#include <common/utils/process_path.h>

#include "KeyboardManagerConstants.h"

namespace Helpers
{
    DWORD EncodeKeyNumpadOrigin(const DWORD key, const bool extended)
    {
        bool numpad_originated = false;
        switch (key)
        {
        case VK_LEFT:
        case VK_RIGHT:
        case VK_UP:
        case VK_DOWN:
        case VK_INSERT:
        case VK_DELETE:
        case VK_PRIOR:
        case VK_NEXT:
        case VK_HOME:
        case VK_END:
            numpad_originated = !extended;
            break;
        case VK_RETURN:
        case VK_DIVIDE:
            numpad_originated = extended;
            break;
        }

        if (numpad_originated)
            return key | GetNumpadOriginEncodingBit();
        else
            return key;
    }

    DWORD ClearKeyNumpadOrigin(const DWORD key)
    {
        return (key & ~GetNumpadOriginEncodingBit());
    }

    bool IsNumpadOriginated(const DWORD key)
    {
        return !!(key & GetNumpadOriginEncodingBit());
    }

    // Check if it's one of the numpad keys that the shift key can affect, so we can introduce a workaround when they are mapped to shift.
    bool IsNumpadKeyThatIsAffectedByShift(const DWORD vkCode)
    {
        switch (vkCode)
        {
        case VK_NUMPAD0:
        case VK_NUMPAD1:
        case VK_NUMPAD2:
        case VK_NUMPAD3:
        case VK_NUMPAD4:
        case VK_NUMPAD5:
        case VK_NUMPAD6:
        case VK_NUMPAD7:
        case VK_NUMPAD8:
        case VK_NUMPAD9:
        case VK_DECIMAL:
            return true;
        }
        return false;
    }

    DWORD GetNumpadOriginEncodingBit()
    {
        // Intentionally do not mimic KF_EXTENDED to avoid confusion, because it's not the same thing
        // See EncodeKeyNumpadOrigin.
        return 1ull << 31;
    }
    // Function to check if the key is a modifier key
    bool IsModifierKey(DWORD key)
    {
        return (GetKeyType(key) != KeyType::Action);
    }

    // Function to get the combined key for modifier keys
    DWORD GetCombinedKey(DWORD key)
    {
        switch (key)
        {
        case VK_LWIN:
        case VK_RWIN:
            return CommonSharedConstants::VK_WIN_BOTH;
        case VK_LCONTROL:
        case VK_RCONTROL:
            return VK_CONTROL;
        case VK_LMENU:
        case VK_RMENU:
            return VK_MENU;
        case VK_LSHIFT:
        case VK_RSHIFT:
            return VK_SHIFT;
        default:
            return key;
        }
    }

    // Function to get the type of the key
    KeyType GetKeyType(DWORD key)
    {
        switch (key)
        {
        case CommonSharedConstants::VK_WIN_BOTH:
        case VK_LWIN:
        case VK_RWIN:
            return KeyType::Win;
        case VK_CONTROL:
        case VK_LCONTROL:
        case VK_RCONTROL:
            return KeyType::Ctrl;
        case VK_MENU:
        case VK_LMENU:
        case VK_RMENU:
            return KeyType::Alt;
        case VK_SHIFT:
        case VK_LSHIFT:
        case VK_RSHIFT:
            return KeyType::Shift;
        default:
            return KeyType::Action;
        }
    }

    // Function to return if the key is an extended key which requires the use of the extended key flag
    bool IsExtendedKey(DWORD key)
    {
        switch (key)
        {
        case VK_RCONTROL:
        case VK_RMENU:
        case VK_NUMLOCK:
        case VK_SNAPSHOT:
        case VK_CANCEL:
            // If the extended flag is not set for the following keys, their NumPad versions are sent. This causes weird behavior when NumLock is on (more information at https://github.com/microsoft/PowerToys/issues/3478)
        case VK_INSERT:
        case VK_HOME:
        case VK_PRIOR:
        case VK_DELETE:
        case VK_END:
        case VK_NEXT:
        case VK_LEFT:
        case VK_DOWN:
        case VK_RIGHT:
        case VK_UP:
        case VK_SLEEP:
        case VK_MEDIA_NEXT_TRACK:
        case VK_MEDIA_PREV_TRACK:
        case VK_MEDIA_STOP:
        case VK_MEDIA_PLAY_PAUSE:
        case VK_VOLUME_MUTE:
        case VK_VOLUME_UP:
        case VK_VOLUME_DOWN:
        case VK_LAUNCH_MEDIA_SELECT:
        case VK_LAUNCH_MAIL:
        case VK_LAUNCH_APP1:
        case VK_LAUNCH_APP2:
        case VK_BROWSER_SEARCH:
        case VK_BROWSER_HOME:
        case VK_BROWSER_BACK:
        case VK_BROWSER_FORWARD:
        case VK_BROWSER_STOP:
        case VK_BROWSER_REFRESH:
        case VK_BROWSER_FAVORITES:
            return true;
        default:
            return false;
        }
    }

    // Function to set the value of a key event based on the arguments
    void SetKeyEvent(std::vector<INPUT>& keyEventArray, DWORD inputType, WORD keyCode, DWORD flags, ULONG_PTR extraInfo)
    {
        INPUT keyEvent{};
        keyEvent.type = inputType;
        keyEvent.ki.wVk = keyCode;
        keyEvent.ki.dwFlags = flags;
        if (IsExtendedKey(keyCode))
        {
            keyEvent.ki.dwFlags |= KEYEVENTF_EXTENDEDKEY;
        }
        keyEvent.ki.dwExtraInfo = extraInfo;

        // Set wScan to the value from MapVirtualKey as some applications may use the scan code for handling input, for instance, Windows Terminal ignores non-character input which has scancode set to 0.
        // MapVirtualKey returns 0 if the key code does not correspond to a physical key (such as unassigned/reserved keys). More details at https://github.com/microsoft/PowerToys/pull/7143#issue-498877747
        keyEvent.ki.wScan = static_cast<WORD>(MapVirtualKey(keyCode, MAPVK_VK_TO_VSC));
        keyEventArray.push_back(keyEvent);
    }

    // Function to set the dummy key events used for remapping shortcuts, required to ensure releasing a modifier doesn't trigger another action (For example, Win->Start Menu or Alt->Menu bar)
    void SetDummyKeyEvent(std::vector<INPUT>& keyEventArray, ULONG_PTR extraInfo)
    {
        SetKeyEvent(keyEventArray, INPUT_KEYBOARD, static_cast<WORD>(KeyboardManagerConstants::DUMMY_KEY), 0, extraInfo);
        SetKeyEvent(keyEventArray, INPUT_KEYBOARD, static_cast<WORD>(KeyboardManagerConstants::DUMMY_KEY), KEYEVENTF_KEYUP, extraInfo);
    }

    // Function to return window handle for a full screen UWP app
    HWND GetFullscreenUWPWindowHandle()
    {
        // Using GetGUIThreadInfo for getting the process of the window in focus. GetForegroundWindow has issues with UWP apps as it returns the Application Frame Host as its linked process
        GUITHREADINFO guiThreadInfo;
        guiThreadInfo.cbSize = sizeof(GUITHREADINFO);
        GetGUIThreadInfo(0, &guiThreadInfo);

        // If no window in focus, use the active window
        if (guiThreadInfo.hwndFocus == nullptr)
        {
            return guiThreadInfo.hwndActive;
        }
        return guiThreadInfo.hwndFocus;
    }

    // Function to return the executable name of the application in focus
    std::wstring GetCurrentApplication(bool keepPath)
    {
        HWND current_window_handle = GetForegroundWindow();
        std::wstring process_name;

        if (current_window_handle != nullptr)
        {
            std::wstring process_path = get_process_path(current_window_handle);
            process_name = process_path;

            // Get process name from path
            PathStripPath(&process_path[0]);

            // Remove elements after null character
            process_path.erase(std::find(process_path.begin(), process_path.end(), L'\0'), process_path.end());

            // If the UWP app is in full-screen, then using GetForegroundWindow approach might fail
            if (process_path == L"ApplicationFrameHost.exe")
            {
                HWND fullscreen_window_handle = GetFullscreenUWPWindowHandle();
                if (fullscreen_window_handle != nullptr)
                {
                    process_path = get_process_path(fullscreen_window_handle);
                    process_name = process_path;

                    // Get process name from path
                    PathStripPath(&process_path[0]);

                    // Remove elements after null character
                    process_path.erase(std::find(process_path.begin(), process_path.end(), L'\0'), process_path.end());
                }
            }

            // If keepPath is false, then return only the name of the process
            if (!keepPath)
            {
                process_name = process_path;
            }
        }

        return process_name;
    }

    // Function to set key events for modifier keys: When shortcutToCompare is passed (non-empty shortcut), then the key event is sent only if both shortcut's don't have the same modifier key. When keyToBeReleased is passed (non-NULL), then the key event is sent if either the shortcuts don't have the same modifier or if the shortcutToBeSent's modifier matches the keyToBeReleased
    void SetModifierKeyEvents(const Shortcut& shortcutToBeSent, const Modifiers& modifiersKeys, std::vector<INPUT>& keyEventArray, bool isKeyDown, ULONG_PTR extraInfoFlag, const Shortcut& shortcutToCompare, const DWORD& keyToBeReleased)
    {
        // If key down is to be sent, send in the order Win, Ctrl, Alt, Shift
        if (isKeyDown)
        {
            // If shortcutToCompare is non-empty, then the key event is sent only if both shortcut's don't have the same modifier key. If keyToBeReleased is non-NULL, then the key event is sent if either the shortcuts don't have the same modifier or if the shortcutToBeSent's modifier matches the keyToBeReleased
            if (shortcutToBeSent.GetWinKey(modifiersKeys.winKey) != NULL && (shortcutToCompare.IsEmpty() || shortcutToBeSent.GetWinKey(modifiersKeys.winKey) != shortcutToCompare.GetWinKey(modifiersKeys.winKey)) && (keyToBeReleased == NULL || !shortcutToBeSent.CheckWinKey(keyToBeReleased)))
            {
                Helpers::SetKeyEvent(keyEventArray, INPUT_KEYBOARD, static_cast<WORD>(shortcutToBeSent.GetWinKey(modifiersKeys.winKey)), 0, extraInfoFlag);
            }
            if (shortcutToBeSent.GetCtrlKey(modifiersKeys.ctrlKey) != NULL && (shortcutToCompare.IsEmpty() || shortcutToBeSent.GetCtrlKey(modifiersKeys.ctrlKey) != shortcutToCompare.GetCtrlKey(modifiersKeys.ctrlKey)) && (keyToBeReleased == NULL || !shortcutToBeSent.CheckCtrlKey(keyToBeReleased)))
            {
                Helpers::SetKeyEvent(keyEventArray, INPUT_KEYBOARD, static_cast<WORD>(shortcutToBeSent.GetCtrlKey(modifiersKeys.ctrlKey)), 0, extraInfoFlag);
            }
            if (shortcutToBeSent.GetAltKey(modifiersKeys.altKey) != NULL && (shortcutToCompare.IsEmpty() || shortcutToBeSent.GetAltKey(modifiersKeys.altKey) != shortcutToCompare.GetAltKey(modifiersKeys.altKey)) && (keyToBeReleased == NULL || !shortcutToBeSent.CheckAltKey(keyToBeReleased)))
            {
                Helpers::SetKeyEvent(keyEventArray, INPUT_KEYBOARD, static_cast<WORD>(shortcutToBeSent.GetAltKey(modifiersKeys.altKey)), 0, extraInfoFlag);
            }
            if (shortcutToBeSent.GetShiftKey(modifiersKeys.shiftKey) != NULL && (shortcutToCompare.IsEmpty() || shortcutToBeSent.GetShiftKey(modifiersKeys.shiftKey) != shortcutToCompare.GetShiftKey(modifiersKeys.shiftKey)) && (keyToBeReleased == NULL || !shortcutToBeSent.CheckShiftKey(keyToBeReleased)))
            {
                Helpers::SetKeyEvent(keyEventArray, INPUT_KEYBOARD, static_cast<WORD>(shortcutToBeSent.GetShiftKey(modifiersKeys.shiftKey)), 0, extraInfoFlag);
            }
        }

        // If key up is to be sent, send in the order Shift, Alt, Ctrl, Win
        else
        {
            // If shortcutToCompare is non-empty, then the key event is sent only if both shortcut's don't have the same modifier key. If keyToBeReleased is non-NULL, then the key event is sent if either the shortcuts don't have the same modifier or if the shortcutToBeSent's modifier matches the keyToBeReleased
            if (shortcutToBeSent.GetShiftKey(modifiersKeys.shiftKey) != NULL && (shortcutToCompare.IsEmpty() || shortcutToBeSent.GetShiftKey(modifiersKeys.shiftKey) != shortcutToCompare.GetShiftKey(modifiersKeys.shiftKey) || shortcutToBeSent.CheckShiftKey(keyToBeReleased)))
            {
                Helpers::SetKeyEvent(keyEventArray, INPUT_KEYBOARD, static_cast<WORD>(shortcutToBeSent.GetShiftKey(modifiersKeys.shiftKey)), KEYEVENTF_KEYUP, extraInfoFlag);
            }
            if (shortcutToBeSent.GetAltKey(modifiersKeys.altKey) != NULL && (shortcutToCompare.IsEmpty() || shortcutToBeSent.GetAltKey(modifiersKeys.altKey) != shortcutToCompare.GetAltKey(modifiersKeys.altKey) || shortcutToBeSent.CheckAltKey(keyToBeReleased)))
            {
                Helpers::SetKeyEvent(keyEventArray, INPUT_KEYBOARD, static_cast<WORD>(shortcutToBeSent.GetAltKey(modifiersKeys.altKey)), KEYEVENTF_KEYUP, extraInfoFlag);
            }
            if (shortcutToBeSent.GetCtrlKey(modifiersKeys.ctrlKey) != NULL && (shortcutToCompare.IsEmpty() || shortcutToBeSent.GetCtrlKey(modifiersKeys.ctrlKey) != shortcutToCompare.GetCtrlKey(modifiersKeys.ctrlKey) || shortcutToBeSent.CheckCtrlKey(keyToBeReleased)))
            {
                Helpers::SetKeyEvent(keyEventArray, INPUT_KEYBOARD, static_cast<WORD>(shortcutToBeSent.GetCtrlKey(modifiersKeys.ctrlKey)), KEYEVENTF_KEYUP, extraInfoFlag);
            }
            if (shortcutToBeSent.GetWinKey(modifiersKeys.winKey) != NULL && (shortcutToCompare.IsEmpty() || shortcutToBeSent.GetWinKey(modifiersKeys.winKey) != shortcutToCompare.GetWinKey(modifiersKeys.winKey) || shortcutToBeSent.CheckWinKey(keyToBeReleased)))
            {
                Helpers::SetKeyEvent(keyEventArray, INPUT_KEYBOARD, static_cast<WORD>(shortcutToBeSent.GetWinKey(modifiersKeys.winKey)), KEYEVENTF_KEYUP, extraInfoFlag);
            }
        }
    }

    // Sends text input directly via SendInput, handling newlines by sending
    // Shift+Enter. Each character is sent individually to avoid a synchronization
    // error across key-down and key-up events that causes repeated or dropped characters
    // when large batches of KEYEVENTF_UNICODE events are sent at once.
    void SendTextInput(const std::wstring& text)
    {
        for (size_t i = 0; i < text.size(); ++i)
        {
            wchar_t c = text[i];

            // Handle \r\n as a single newline
            if (c == L'\r' && i + 1 < text.size() && text[i + 1] == L'\n')
            {
                ++i;
            }

            if (c == L'\r' || c == L'\n')
            {
                // Send Shift+Enter instead of bare Enter so that chat apps
                // (Teams, Slack, Discord, etc.) insert a new line rather than
                // submitting the message. In plain text editors both behave
                // the same.
                INPUT returnInputs[4]{};

                // Shift down
                returnInputs[0].type = INPUT_KEYBOARD;
                returnInputs[0].ki.wVk = VK_SHIFT;
                returnInputs[0].ki.wScan = static_cast<WORD>(MapVirtualKey(VK_SHIFT, MAPVK_VK_TO_VSC));
                returnInputs[0].ki.dwExtraInfo = KeyboardManagerConstants::KEYBOARDMANAGER_SHORTCUT_FLAG;

                // Return down
                returnInputs[1].type = INPUT_KEYBOARD;
                returnInputs[1].ki.wVk = VK_RETURN;
                returnInputs[1].ki.wScan = static_cast<WORD>(MapVirtualKey(VK_RETURN, MAPVK_VK_TO_VSC));
                returnInputs[1].ki.dwExtraInfo = KeyboardManagerConstants::KEYBOARDMANAGER_SHORTCUT_FLAG;

                // Return up
                returnInputs[2].type = INPUT_KEYBOARD;
                returnInputs[2].ki.wVk = VK_RETURN;
                returnInputs[2].ki.dwFlags = KEYEVENTF_KEYUP;
                returnInputs[2].ki.wScan = static_cast<WORD>(MapVirtualKey(VK_RETURN, MAPVK_VK_TO_VSC));
                returnInputs[2].ki.dwExtraInfo = KeyboardManagerConstants::KEYBOARDMANAGER_SHORTCUT_FLAG;

                // Shift up
                returnInputs[3].type = INPUT_KEYBOARD;
                returnInputs[3].ki.wVk = VK_SHIFT;
                returnInputs[3].ki.dwFlags = KEYEVENTF_KEYUP;
                returnInputs[3].ki.wScan = static_cast<WORD>(MapVirtualKey(VK_SHIFT, MAPVK_VK_TO_VSC));
                returnInputs[3].ki.dwExtraInfo = KeyboardManagerConstants::KEYBOARDMANAGER_SHORTCUT_FLAG;

                SendInput(ARRAYSIZE(returnInputs), returnInputs, sizeof(INPUT));
                continue;
            }

            INPUT charInputs[2]{};
            charInputs[0].type = INPUT_KEYBOARD;
            charInputs[0].ki.dwFlags = KEYEVENTF_UNICODE;
            charInputs[0].ki.dwExtraInfo = KeyboardManagerConstants::KEYBOARDMANAGER_SHORTCUT_FLAG;
            charInputs[0].ki.wScan = c;

            charInputs[1].type = INPUT_KEYBOARD;
            charInputs[1].ki.dwFlags = KEYEVENTF_UNICODE | KEYEVENTF_KEYUP;
            charInputs[1].ki.dwExtraInfo = KeyboardManagerConstants::KEYBOARDMANAGER_SHORTCUT_FLAG;
            charInputs[1].ki.wScan = c;

            SendInput(ARRAYSIZE(charInputs), charInputs, sizeof(INPUT));
        }
    }

    // Function to filter the key codes for artificial key codes
    int32_t FilterArtificialKeys(const int32_t& key)
    {
        switch (key)
        {
        case CommonSharedConstants::VK_WIN_BOTH:
            return VK_LWIN;
        }

        return key;
    }

    // Function to sort a vector of shortcuts based on its size
    void SortShortcutVectorBasedOnSize(std::vector<Shortcut>& shortcutVector)
    {
        std::sort(shortcutVector.begin(), shortcutVector.end(), [](Shortcut first, Shortcut second) {
            return first.Size() > second.Size();
        });
    }
}
