#include <windows.h>
#include <shellapi.h>
#include <commctrl.h>
#include "resource.h"

#pragma comment(lib, "comctl32.lib")
#pragma comment(lib, "user32.lib")
#pragma comment(lib, "shell32.lib")

#pragma comment(linker, \
    "\"/manifestdependency:type='win32' " \
    "name='Microsoft.Windows.Common-Controls' " \
    "version='6.0.0.0' " \
    "processorArchitecture='*' " \
    "publicKeyToken='6595b64144ccf1df' " \
    "language='*'\"")

#define WM_TRAYNOTIFY   (WM_USER + 100)
#define ID_TRAYICON     1
#define IDM_EXIT        1001
#define IDM_TOGGLE      1002
#define IDM_ABOUT       1003

HHOOK           hHook = NULL;
HWND            hWndHidden = NULL;
NOTIFYICONDATA  nid = { 0 };
BOOL            enabled = TRUE;
BOOL            keystrokeRightAltProcessed = FALSE;

HINSTANCE       g_hInstance = NULL;

LPCTSTR LoadStr(UINT id)
{
    static TCHAR buf[1024];
    buf[0] = TEXT('\0');
    LoadString(g_hInstance, id, buf, _countof(buf));
    return buf;
}

void ShowError(UINT id)
{
    MessageBox(NULL, LoadStr(id), TEXT("SwitchyAlt"), MB_OK | MB_ICONERROR);
}

void SwitchLayout()
{
    HWND hwndForeground = GetForegroundWindow();
    if (!hwndForeground) return;

    char className[256];
    GetClassNameA(hwndForeground, className, sizeof(className));

    BOOL isRDP = (strstr(className, "TscShellContainerClass") != NULL);
    BOOL isConsole = (strstr(className, "ConsoleWindowClass") != NULL);

    if (isRDP || isConsole)
    {
        INPUT inputs[4] = { 0 };
        for (int i = 0; i < 4; i++) inputs[i].type = INPUT_KEYBOARD;
        inputs[0].ki.wVk = VK_MENU;  inputs[0].ki.wScan = 0x38;
        inputs[1].ki.wVk = VK_SHIFT; inputs[1].ki.wScan = 0x2A;
        inputs[2].ki.wVk = VK_SHIFT; inputs[2].ki.wScan = 0x2A; inputs[2].ki.dwFlags = KEYEVENTF_KEYUP;
        inputs[3].ki.wVk = VK_MENU;  inputs[3].ki.wScan = 0x38; inputs[3].ki.dwFlags = KEYEVENTF_KEYUP;
        SendInput(4, inputs, sizeof(INPUT));
    }
    else
    {
        GUITHREADINFO gti = { sizeof(GUITHREADINFO) };
        HWND targetWnd = hwndForeground;

        if (GetGUIThreadInfo(0, &gti) && gti.hwndFocus != NULL)
        {
            targetWnd = gti.hwndFocus;
        }

        PostMessage(targetWnd, WM_INPUTLANGCHANGEREQUEST, 0, (LPARAM)HKL_NEXT);
    }
}

LRESULT CALLBACK LowLevelKeyboardProc(int nCode, WPARAM wParam, LPARAM lParam)
{
    if (nCode != HC_ACTION)
        return CallNextHookEx(hHook, nCode, wParam, lParam);

    KBDLLHOOKSTRUCT* p = (KBDLLHOOKSTRUCT*)lParam;

    if (p->flags & LLKHF_INJECTED)
        return CallNextHookEx(hHook, nCode, wParam, lParam);

    if (p->vkCode == VK_RMENU)
    {
        if (wParam == WM_KEYDOWN || wParam == WM_SYSKEYDOWN)
        {
            if (!keystrokeRightAltProcessed)
            {
                keystrokeRightAltProcessed = TRUE;

                if (GetAsyncKeyState(VK_LMENU) & 0x8000)
                {
                    enabled = !enabled;
                    return 1;
                }

                if (enabled)
                {
                    SwitchLayout();
                }
            }
            return 1;
        }

        if (wParam == WM_KEYUP || wParam == WM_SYSKEYUP)
        {
            keystrokeRightAltProcessed = FALSE;
            return 1;
        }
    }

    return CallNextHookEx(hHook, nCode, wParam, lParam);
}

INT_PTR CALLBACK AboutDlgProc(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam)
{
    switch (message)
    {
    case WM_INITDIALOG:
        SetDlgItemText(hDlg, IDC_INFO, LoadStr(IDS_ABOUT_TEXT));
        SetDlgItemText(hDlg, IDC_STATIC_APPINFO, LoadStr(IDS_APPINFO));
        SetDlgItemText(hDlg, IDC_STATIC_BASEDON, LoadStr(IDS_BASEDON));
        return (INT_PTR)TRUE;

    case WM_COMMAND:
        if (LOWORD(wParam) == IDOK || LOWORD(wParam) == IDCANCEL)
        {
            EndDialog(hDlg, LOWORD(wParam));
            return (INT_PTR)TRUE;
        }
        break;
    }
    return (INT_PTR)FALSE;
}

LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam)
{
    switch (msg)
    {
    case WM_TRAYNOTIFY:
        if (lParam == WM_RBUTTONUP)
        {
            POINT pt;
            GetCursorPos(&pt);

            HMENU hMenu = CreatePopupMenu();
            if (hMenu)
            {
                AppendMenu(hMenu, MF_STRING, IDM_TOGGLE,
                    enabled ? LoadStr(IDS_MENU_PAUSE) : LoadStr(IDS_MENU_RESUME));

                AppendMenu(hMenu, MF_STRING, IDM_ABOUT, LoadStr(IDS_MENU_ABOUT));
                AppendMenu(hMenu, MF_SEPARATOR, 0, NULL);
                AppendMenu(hMenu, MF_STRING, IDM_EXIT, LoadStr(IDS_MENU_EXIT));

                SetForegroundWindow(hwnd);
                UINT cmd = TrackPopupMenu(hMenu, TPM_RETURNCMD | TPM_NONOTIFY, pt.x, pt.y, 0, hwnd, NULL);
                DestroyMenu(hMenu);

                if (cmd == IDM_TOGGLE) enabled = !enabled;
                if (cmd == IDM_ABOUT)
                    DialogBox(g_hInstance, MAKEINTRESOURCE(IDD_ABOUTBOX), hwnd, AboutDlgProc);
                if (cmd == IDM_EXIT)
                    PostQuitMessage(0);
            }
        }
        break;

    case WM_DESTROY:
        PostQuitMessage(0);
        break;

    default:
        return DefWindowProc(hwnd, msg, wParam, lParam);
    }
    return 0;
}

int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrev, LPSTR lpCmd, int nShow)
{
    g_hInstance = hInstance;

    HANDLE hMutex = CreateMutex(NULL, FALSE, TEXT("SwitchyAltUnique"));
    if (GetLastError() == ERROR_ALREADY_EXISTS)
    {
        ShowError(IDS_APP_MUTEX_ERROR);
        return 1;
    }

    WNDCLASSEX wc = { sizeof(WNDCLASSEX) };
    wc.lpfnWndProc = WndProc;
    wc.hInstance = hInstance;
    wc.lpszClassName = TEXT("SwitchyHiddenClass");

    if (!RegisterClassEx(&wc))
    {
        ShowError(IDS_ERROR_REG_CLASS);
        return 1;
    }

    hWndHidden = CreateWindowEx(0, wc.lpszClassName, TEXT("SwitchyAlt"),
        0, 0, 0, 0, 0, HWND_MESSAGE, NULL, hInstance, NULL);

    if (!hWndHidden)
    {
        ShowError(IDS_ERROR_CREATE_WINDOW);
        return 1;
    }

    InitCommonControls();

    nid.cbSize = sizeof(NOTIFYICONDATA);
    nid.hWnd = hWndHidden;
    nid.uID = ID_TRAYICON;
    nid.uFlags = NIF_ICON | NIF_MESSAGE | NIF_TIP;
    nid.uCallbackMessage = WM_TRAYNOTIFY;
    nid.hIcon = LoadIcon(hInstance, MAKEINTRESOURCE(IDI_ICON1));
    lstrcpy(nid.szTip, LoadStr(IDS_TRAY_TOOLTIP));

    if (!Shell_NotifyIcon(NIM_ADD, &nid))
    {
        ShowError(IDS_ERROR_TRAY_ICON);
    }

    hHook = SetWindowsHookEx(WH_KEYBOARD_LL, LowLevelKeyboardProc, hInstance, 0);
    if (!hHook)
    {
        ShowError(IDS_ERROR_SETHOOK);
        Shell_NotifyIcon(NIM_DELETE, &nid);
        DestroyWindow(hWndHidden);
        return 1;
    }

    MSG msg;
    while (GetMessage(&msg, NULL, 0, 0))
    {
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }

    UnhookWindowsHookEx(hHook);
    Shell_NotifyIcon(NIM_DELETE, &nid);
    DestroyWindow(hWndHidden);

    return (int)msg.wParam;
}