#include <windows.h>
#include <shellapi.h>
#include <wtsapi32.h>
#include <cwchar>
#include <initializer_list>
#include <string>

namespace {
#ifdef STAYHERE_TEST_INSTANCE
constexpr wchar_t ClassName[] = L"StayHere.Test.Tray.1";
constexpr wchar_t MutexName[] = L"Local\\StayHere.Test.Tray.1";
#else
constexpr wchar_t ClassName[] = L"StayHere.Tray.1";
constexpr wchar_t MutexName[] = L"Local\\StayHere.Tray.1";
#endif
constexpr UINT TrayMessage = WM_APP + 1;
constexpr UINT ControlMessage = WM_APP + 2;
constexpr UINT Toggle = 1, Quit = 2, Pause = 3, Resume = 4, Status = 5;
constexpr DWORD IdleThreshold = 180000;
bool enabled = true;
bool sessionLocked = false;
bool powerOk = true;
bool inputOk = true;
UINT taskbarCreated;
HICON activeIcon, pausedIcon;
NOTIFYICONDATAW tray{};

constexpr bool shouldSignal(bool active, bool unlocked, DWORD now, DWORD lastInput) {
    return active && unlocked && DWORD(now - lastInput) >= IdleThreshold;
}
static_assert(!shouldSignal(true, true, 179999, 0));
static_assert(shouldSignal(true, true, 180000, 0));
static_assert(!shouldSignal(false, true, 180000, 0));
static_assert(!shouldSignal(true, false, 180000, 0));
static_assert(shouldSignal(true, true, 179000, 0xfffffc18)); // tick counter wrap

bool interactiveDesktop() {
    if (sessionLocked) return false;
    HDESK desktop = OpenInputDesktop(0, FALSE, DESKTOP_READOBJECTS);
    if (!desktop) return false;
    wchar_t name[128]{};
    DWORD needed = 0;
    bool ok = GetUserObjectInformationW(desktop, UOI_NAME, name, sizeof(name), &needed)
        && _wcsicmp(name, L"Default") == 0;
    CloseDesktop(desktop);
    if (!ok) return false;
    DWORD session;
    LPWSTR buffer = nullptr;
    DWORD bytes = 0;
    if (!ProcessIdToSessionId(GetCurrentProcessId(), &session)) return false;
    if (!WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, session, WTSConnectState, &buffer, &bytes)) return false;
    ok = bytes >= sizeof(WTS_CONNECTSTATE_CLASS)
        && *reinterpret_cast<WTS_CONNECTSTATE_CLASS*>(buffer) == WTSActive;
    WTSFreeMemory(buffer);
    return ok;
}

HICON makeIcon(COLORREF color) {
    HDC screen = GetDC(nullptr);
    HDC dc = CreateCompatibleDC(screen);
    HBITMAP bitmap = CreateCompatibleBitmap(screen, 32, 32);
    HGDIOBJ previous = SelectObject(dc, bitmap);
    RECT rect{0, 0, 32, 32};
    FillRect(dc, &rect, static_cast<HBRUSH>(GetStockObject(BLACK_BRUSH)));
    HBRUSH brush = CreateSolidBrush(color);
    HGDIOBJ oldBrush = SelectObject(dc, brush);
    HGDIOBJ oldPen = SelectObject(dc, GetStockObject(NULL_PEN));
    Ellipse(dc, 3, 3, 29, 29);
    SelectObject(dc, oldBrush);
    SelectObject(dc, oldPen);
    DeleteObject(brush);
    SelectObject(dc, previous);
    HBITMAP mask = CreateBitmap(32, 32, 1, 1, nullptr);
    HGDIOBJ maskPrevious = SelectObject(dc, mask);
    PatBlt(dc, 0, 0, 32, 32, WHITENESS);
    oldBrush = SelectObject(dc, GetStockObject(BLACK_BRUSH));
    oldPen = SelectObject(dc, GetStockObject(NULL_PEN));
    Ellipse(dc, 3, 3, 29, 29);
    SelectObject(dc, oldPen);
    SelectObject(dc, oldBrush);
    SelectObject(dc, maskPrevious);
    ICONINFO info{TRUE, 0, 0, mask, bitmap};
    HICON icon = CreateIconIndirect(&info);
    DeleteObject(bitmap);
    DeleteObject(mask);
    DeleteDC(dc);
    ReleaseDC(nullptr, screen);
    return icon;
}

void updateTray() {
    tray.hIcon = enabled ? activeIcon : pausedIcon;
    const wchar_t* tip = !enabled ? L"StayHere — призупинено"
        : !powerOk ? L"StayHere — помилка утримання від сну"
        : !inputOk ? L"StayHere — не вдалося надіслати F15"
        : L"StayHere — активно (F15 після 3 хв простою)";
    lstrcpynW(tray.szTip, tip, ARRAYSIZE(tray.szTip));
    Shell_NotifyIconW(NIM_MODIFY, &tray);
}

void setEnabled(bool value) {
    enabled = value;
    inputOk = true;
    powerOk = SetThreadExecutionState(ES_CONTINUOUS |
        (enabled && interactiveDesktop() ? ES_SYSTEM_REQUIRED | ES_DISPLAY_REQUIRED : 0)) != 0;
    updateTray();
}

void tick() {
    const bool unlocked = interactiveDesktop();
    powerOk = SetThreadExecutionState(ES_CONTINUOUS |
        (enabled && unlocked ? ES_SYSTEM_REQUIRED | ES_DISPLAY_REQUIRED : 0)) != 0;
    LASTINPUTINFO idle{sizeof(idle), 0};
    if (GetLastInputInfo(&idle) && shouldSignal(enabled, unlocked, GetTickCount(), idle.dwTime)) {
        // Avoid inserting a key into a chord being held while the user is away.
        for (int key : {VK_SHIFT, VK_CONTROL, VK_MENU, VK_LWIN, VK_RWIN,
                        VK_LBUTTON, VK_RBUTTON, VK_MBUTTON}) {
            if (GetAsyncKeyState(key) & 0x8000) return;
        }
        INPUT input[2]{};
        input[0].type = input[1].type = INPUT_KEYBOARD;
        input[0].ki.wVk = input[1].ki.wVk = VK_F15;
        input[1].ki.dwFlags = KEYEVENTF_KEYUP;
        inputOk = SendInput(2, input, sizeof(INPUT)) == 2;
        if (!inputOk) { // If only key-down succeeded, make a best effort to release it.
            SendInput(1, &input[1], sizeof(INPUT));
        }
    }
    updateTray();
}

void showMenu(HWND window) {
    HMENU menu = CreatePopupMenu();
    AppendMenuW(menu, MF_STRING | (enabled ? MF_CHECKED : 0), Toggle, L"Підтримувати активність");
    AppendMenuW(menu, MF_SEPARATOR, 0, nullptr);
    AppendMenuW(menu, MF_STRING, Quit, L"Вийти");
    POINT pos;
    GetCursorPos(&pos);
    SetForegroundWindow(window);
    UINT command = TrackPopupMenu(menu, TPM_RETURNCMD | TPM_RIGHTBUTTON, pos.x, pos.y, 0, window, nullptr);
    DestroyMenu(menu);
    if (command) SendMessageW(window, ControlMessage, command, 0);
    PostMessageW(window, WM_NULL, 0, 0);
}

LRESULT CALLBACK windowProc(HWND window, UINT message, WPARAM wp, LPARAM lp) {
    if (message == taskbarCreated) {
        Shell_NotifyIconW(NIM_ADD, &tray);
        return 0;
    }
    switch (message) {
    case TrayMessage:
        if (lp == WM_RBUTTONUP || lp == WM_CONTEXTMENU) showMenu(window);
        else if (lp == WM_LBUTTONUP) setEnabled(!enabled);
        return 0;
    case ControlMessage:
        if (wp == Status) return enabled ? 1 : 2;
        if (wp == Toggle) setEnabled(!enabled);
        else if (wp == Pause) setEnabled(false);
        else if (wp == Resume) setEnabled(true);
        else if (wp == Quit) DestroyWindow(window);
        return 0;
    case WM_TIMER: tick(); return 0;
    case WM_WTSSESSION_CHANGE:
        if (wp == WTS_SESSION_LOCK || wp == WTS_REMOTE_DISCONNECT || wp == WTS_CONSOLE_DISCONNECT)
            sessionLocked = true;
        else if (wp == WTS_SESSION_UNLOCK || wp == WTS_REMOTE_CONNECT || wp == WTS_CONSOLE_CONNECT)
            sessionLocked = false;
        setEnabled(enabled);
        return 0;
    case WM_DESTROY:
        KillTimer(window, 1);
        WTSUnRegisterSessionNotification(window);
        SetThreadExecutionState(ES_CONTINUOUS);
        Shell_NotifyIconW(NIM_DELETE, &tray);
        PostQuitMessage(0);
        return 0;
    }
    return DefWindowProcW(window, message, wp, lp);
}
}

int WINAPI wWinMain(HINSTANCE instance, HINSTANCE, PWSTR, int) {
    int argc{};
    auto argv=CommandLineToArgvW(GetCommandLineW(),&argc);
    if(!argv)return 10;
    const std::wstring argument=argc==2?argv[1]:L"";
    LocalFree(argv);
    if(argc>2)return 10;
    const auto commandLine=argument.c_str();
    UINT command = 0;
    if (wcscmp(commandLine, L"--pause") == 0) command = Pause;
    else if (wcscmp(commandLine, L"--resume") == 0) command = Resume;
    else if (wcscmp(commandLine, L"--exit") == 0) command = Quit;
    else if (wcscmp(commandLine, L"--status") == 0) command = Status;
    else if (*commandLine) return 10;
    HWND existing = FindWindowW(ClassName, nullptr);
    if (command) {
        if (!existing) return 3;
        DWORD_PTR result = 0;
        if (!SendMessageTimeoutW(existing, ControlMessage, command, 0, SMTO_ABORTIFHUNG, 2000, &result)) return 4;
        return command == Status ? static_cast<int>(result) : 0;
    }
    HANDLE mutex = CreateMutexW(nullptr, FALSE, MutexName);
    if (!mutex) return 5;
    if (GetLastError() == ERROR_ALREADY_EXISTS) { CloseHandle(mutex); return 0; }
    taskbarCreated = RegisterWindowMessageW(L"TaskbarCreated");
    activeIcon = makeIcon(RGB(40, 185, 100));
    pausedIcon = makeIcon(RGB(135, 140, 150));
    WNDCLASSW wc{};
    wc.lpfnWndProc = windowProc;
    wc.hInstance = instance;
    wc.lpszClassName = ClassName;
    if (!RegisterClassW(&wc)) return 6;
    HWND window = CreateWindowExW(0, ClassName, L"StayHere", WS_OVERLAPPED,
        0, 0, 0, 0, nullptr, nullptr, instance, nullptr);
    if (!window) return 7;
    tray.cbSize = sizeof(tray);
    tray.hWnd = window;
    tray.uID = 1;
    tray.uFlags = NIF_ICON | NIF_MESSAGE | NIF_TIP;
    tray.uCallbackMessage = TrayMessage;
    tray.hIcon = activeIcon;
    lstrcpyW(tray.szTip, L"StayHere");
    if (!Shell_NotifyIconW(NIM_ADD, &tray)) { DestroyWindow(window); return 8; }
    WTSRegisterSessionNotification(window, NOTIFY_FOR_THIS_SESSION);
    setEnabled(true);
    if (!SetTimer(window, 1, 15000, nullptr)) { DestroyWindow(window); return 9; }
    MSG msg{};
    while (GetMessageW(&msg, nullptr, 0, 0) > 0) {
        TranslateMessage(&msg);
        DispatchMessageW(&msg);
    }
    DestroyIcon(activeIcon);
    DestroyIcon(pausedIcon);
    CloseHandle(mutex);
    return 0;
}