VirtualDesktopIndicator.py

"""
Virtual Desktop Indicator

Shows the current Windows virtual desktop number in a small overlay.
Includes a system tray icon with a menu item to exit the application.

License
-------
MIT

Release Notes
-------------
1.0.0 - Initial release
- Shows the current Windows virtual desktop number in a small overlay
- Supports a transparent click-through display
- Includes a system tray icon with an Exit option
- Can be packaged as a standalone Windows EXE with PyInstaller

Requirements
------------
- Windows 10 or Windows 11
- Python 3.12+
- tkinter (included with the standard Windows Python installer)
- pystray (required for the system tray icon)
- pillow (required by pystray for tray icon image creation)

Install runtime dependencies
----------------------------
Choose one of these approaches:

Option A: install for the current user only
This avoids writing to the global Python installation and is often the safest
choice on managed or locked-down machines.

    py -3.12 -m pip install --user pystray pillow

Option B: install globally
This makes the packages available to all users of this Python installation.
Depending on how Python was installed, you may need to run PowerShell as
Administrator.

    py -3.12 -m pip install pystray pillow

Optional: install PyInstaller for EXE builds
--------------------------------------------
User-only install:

    py -3.12 -m pip install --user pyinstaller

Global install:

    py -3.12 -m pip install pyinstaller

Run from source
---------------
From PowerShell in the folder containing this file:

    py -3.12 .\\VirtualDesktopIndicator.py

Build a single EXE with PyInstaller
-----------------------------------
If PyInstaller was installed for the current user only, or if you want the
most reliable invocation, run it as a Python module:

    py -3.12 -m PyInstaller --onefile --windowed --name VirtualDesktopIndicator .\\VirtualDesktopIndicator.py

If PyInstaller was installed globally and its command is on PATH, this also
works:

    pyinstaller --onefile --windowed --name VirtualDesktopIndicator .\\VirtualDesktopIndicator.py

Build output
------------
The EXE will be created here:

    .\\dist\\VirtualDesktopIndicator.exe

Notes
-----
- --onefile creates a single EXE.
- --windowed prevents a console window from opening.
- The target machine does not need Python installed.
- If Windows Defender or startup speed gets grumpy, build without --onefile:

    py -3.12 -m PyInstaller --windowed --name VirtualDesktopIndicator .\\VirtualDesktopIndicator.py

  That creates a folder-based app under:

    .\\dist\\VirtualDesktopIndicator\\

If pip install fails globally
-----------------------------
Fallback to user-profile installs and run tools as Python modules:

    py -3.12 -m pip install --user pystray pillow pyinstaller
    py -3.12 -m PyInstaller --onefile --windowed --name VirtualDesktopIndicator .\\VirtualDesktopIndicator.py
"""


# =========================
# Metadata
# =========================
__title__ = "Virtual Desktop Indicator"
__license__ = "MIT"
__version__ = "1.0.0"


# =========================
# Imports
# =========================
import ctypes
import ctypes.wintypes as wt
import pystray
from PIL import Image, ImageDraw
import re
import threading
import time
import tkinter as tk
import uuid
import winreg


# =========================
# Config
# =========================
# How often to check whether the active virtual desktop changed, in milliseconds.
POLL_MS = 1000
# Reserved for optional fade-out behavior; currently not used.
DISPLAY_MS = 5000
# Reserved for optional fade animations; currently not used.
FADE_STEP = 0.08
# How often the Tk event loop runs the periodic update callback.
FADE_INTERVAL_MS = 30

# Corner of the primary screen where the overlay should appear.
CORNER = "bottom-left"
# Horizontal distance from the chosen screen edge.
MARGIN_X = 0
# Vertical distance from the chosen screen edge.
MARGIN_Y = 40

# Font used for the desktop number overlay.
FONT = ("Segoe UI", 36, "bold")
# Text color for the desktop number.
TEXT_COLOR = "white"
# Transparency key color. Anything painted with this color becomes invisible.
BG_COLOR = "#ff00ff"
# Window alpha when visible.
BG_ALPHA_VISIBLE = 1.0
# Window alpha when hidden.
BG_ALPHA_HIDDEN = 0.0

# When True, print diagnostic output to the console.
DEBUG = False
# When True, make the overlay ignore mouse clicks.
ENABLE_CLICK_THROUGH = True


# =========================
# Win32 API constants and function prototypes
# =========================
user32 = ctypes.WinDLL("user32", use_last_error=True)
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)

GWL_EXSTYLE = -20
WS_EX_LAYERED = 0x00080000
WS_EX_TRANSPARENT = 0x00000020
WS_EX_TOOLWINDOW = 0x00000080
WS_EX_NOACTIVATE = 0x08000000

SWP_NOMOVE = 0x0002
SWP_NOSIZE = 0x0001
SWP_NOACTIVATE = 0x0010
SWP_FRAMECHANGED = 0x0020
HWND_TOPMOST = -1

user32.GetWindowLongW.argtypes = [wt.HWND, ctypes.c_int]
user32.GetWindowLongW.restype = ctypes.c_long

user32.SetWindowLongW.argtypes = [wt.HWND, ctypes.c_int, ctypes.c_long]
user32.SetWindowLongW.restype = ctypes.c_long

user32.SetWindowPos.argtypes = [
    wt.HWND, wt.HWND, ctypes.c_int, ctypes.c_int,
    ctypes.c_int, ctypes.c_int, ctypes.c_uint
]
user32.SetWindowPos.restype = wt.BOOL

kernel32.GetCurrentProcessId.argtypes = []
kernel32.GetCurrentProcessId.restype = wt.DWORD

kernel32.ProcessIdToSessionId.argtypes = [wt.DWORD, ctypes.POINTER(wt.DWORD)]
kernel32.ProcessIdToSessionId.restype = wt.BOOL


# =========================
# Main Application Class
# =========================
class Overlay:
    def __init__(self):
        """Create the overlay window, position it, and start the tray icon."""
        self.root = tk.Tk()
        self.root.overrideredirect(True)
        self.root.attributes("-topmost", True)
        # self.root.attributes("-alpha", BG_ALPHA_VISIBLE)
        self.root.configure(bg=BG_COLOR)
        self.root.wm_attributes("-transparentcolor", BG_COLOR)

        self.label = tk.Label(
            self.root,
            text="starting...",
            font=FONT,
            fg=TEXT_COLOR,
            bg=BG_COLOR,
            padx=20,
            pady=12,
            bd=0,
            highlightthickness=0,
        )
        self.label.pack()

        self.root.update_idletasks()
        self.root.attributes("-alpha", 1.0)

        self.current_number = None
        self.last_poll = 0.0
        self.tray_icon = None

        self.place_window()

        if ENABLE_CLICK_THROUGH:
            self.root.after(100, self.enable_click_through)

        self.start_tray_icon()

    def create_tray_image(self):
        """Create a simple in-memory tray icon image."""
        image = Image.new("RGBA", (64, 64), (0, 0, 0, 0))
        draw = ImageDraw.Draw(image)
        draw.ellipse((8, 8, 56, 56), fill=(0, 120, 215, 255))
        draw.text((22, 16), "VD", fill=(255, 255, 255, 255))
        return image

    def start_tray_icon(self):
        """Start the system tray icon on a background thread."""
        menu = pystray.Menu(
            pystray.MenuItem("Exit", self.on_tray_exit)
        )
        self.tray_icon = pystray.Icon(
            "VirtualDesktopIndicator",
            self.create_tray_image(),
            "Virtual Desktop Indicator",
            menu,
        )
        threading.Thread(target=self.tray_icon.run, daemon=True).start()

    def on_tray_exit(self, icon, item):
        """Handle the tray menu Exit command on Tk's main thread."""
        self.root.after(0, self.exit_app)

    def exit_app(self):
        """Stop the tray icon and close the overlay cleanly."""
        if self.tray_icon is not None:
            try:
                self.tray_icon.stop()
            except Exception:
                pass
            self.tray_icon = None

        self.root.quit()
        self.root.destroy()

    def enable_click_through(self):
        """Enable mouse click-through mode for the overlay window."""
        hwnd = self.root.winfo_id()
        make_click_through(hwnd)

    def place_window(self):
        """Position the overlay in the configured screen corner."""
        self.root.update_idletasks()
        w = self.root.winfo_reqwidth()
        h = self.root.winfo_reqheight()
        sw = self.root.winfo_screenwidth()
        sh = self.root.winfo_screenheight()

        if CORNER == "top-left":
            x = MARGIN_X
            y = MARGIN_Y
        elif CORNER == "bottom-left":
            x = MARGIN_X
            y = sh - h - MARGIN_Y
        elif CORNER == "bottom-right":
            x = sw - w - MARGIN_X
            y = sh - h - MARGIN_Y
        else:
            x = sw - w - MARGIN_X
            y = MARGIN_Y

        log(f"geometry: {w}x{h}+{x}+{y}")
        self.root.geometry(f"{w}x{h}+{x}+{y}")

    def show_number(self, number):
        """Update the overlay text to the current desktop number."""
        text = str(number) if number is not None else "?"
        log("showing:", text)
        self.label.config(text=text)
        self.place_window()
        self.root.attributes("-alpha", BG_ALPHA_VISIBLE)

    def tick(self):
        """Poll for desktop changes and reschedule the next update."""
        now = time.time()
        if now - self.last_poll >= (POLL_MS / 1000.0):
            self.last_poll = now
            number = get_current_desktop_number()
            if number != self.current_number:
                self.current_number = number
                self.show_number(number)

        self.root.after(FADE_INTERVAL_MS, self.tick)

    def run(self):
        """Show the initial desktop number and start Tk's event loop."""
        self.show_number(get_current_desktop_number())
        self.tick()
        self.root.mainloop()


# =========================
# Helper functions
# =========================
def log(*args):
    """Print debug output only when debugging is enabled."""
    if DEBUG:
        print(*args)

def get_session_id():
    """Return the current Windows session ID for the running process."""
    pid = kernel32.GetCurrentProcessId()
    session_id = wt.DWORD()
    ok = kernel32.ProcessIdToSessionId(pid, ctypes.byref(session_id))
    if not ok:
        raise ctypes.WinError(ctypes.get_last_error())
    return session_id.value

def make_click_through(hwnd):
    """Apply window styles that keep the overlay topmost and mouse-transparent."""
    exstyle = user32.GetWindowLongW(hwnd, GWL_EXSTYLE)
    exstyle |= (
        WS_EX_TRANSPARENT
        | WS_EX_TOOLWINDOW
        | WS_EX_NOACTIVATE
    )
    user32.SetWindowLongW(hwnd, GWL_EXSTYLE, exstyle)
    user32.SetWindowPos(
        hwnd,
        HWND_TOPMOST,
        0,
        0,
        0,
        0,
        SWP_NOMOVE | SWP_NOSIZE | SWP_NOACTIVATE | SWP_FRAMECHANGED,
    )

def read_reg_value(root, path, name):
    """Read a registry value and return None if the key or value is missing."""
    try:
        with winreg.OpenKey(root, path) as key:
            value, _ = winreg.QueryValueEx(key, name)
            return value
    except FileNotFoundError:
        return None
    except OSError:
        return None

def guid_from_bytes_le(b):
    """Convert a little-endian 16-byte GUID buffer into a normalized string."""
    if len(b) != 16:
        return None
    return str(uuid.UUID(bytes_le=b)).lower()

def parse_guidish_string(s):
    """Parse a GUID from a registry string in dashed or plain-hex form."""
    s = s.strip().lower()
    try:
        return str(uuid.UUID(s)).lower()
    except Exception:
        pass

    hex_only = re.sub(r"[^0-9a-f]", "", s)
    if len(hex_only) == 32:
        try:
            return str(uuid.UUID(hex=hex_only)).lower()
        except Exception:
            return None
    return None

def parse_guid_list(value):
    """Parse one or more GUIDs from a registry value containing desktop IDs."""
    if value is None:
        return []

    if isinstance(value, (bytes, bytearray)):
        if len(value) % 16 != 0:
            return []
        out = []
        for i in range(0, len(value), 16):
            g = guid_from_bytes_le(value[i:i + 16])
            if g:
                out.append(g)
        return out

    if isinstance(value, str):
        s = value.strip()

        matches = re.findall(
            r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}",
            s
        )
        if matches:
            return [str(uuid.UUID(m)).lower() for m in matches]

        hex_only = re.sub(r"[^0-9a-fA-F]", "", s)
        if len(hex_only) >= 32 and len(hex_only) % 32 == 0:
            out = []
            for i in range(0, len(hex_only), 32):
                chunk = hex_only[i:i + 32]
                try:
                    out.append(str(uuid.UUID(hex=chunk)).lower())
                except Exception:
                    pass
            return out

    return []

def read_current_desktop_guid():
    """Read the GUID of the currently active virtual desktop from Explorer state."""
    session_id = get_session_id()

    candidate_paths = [
        fr"Software\Microsoft\Windows\CurrentVersion\Explorer\SessionInfo\{session_id}\VirtualDesktops",
        r"Software\Microsoft\Windows\CurrentVersion\Explorer\VirtualDesktops",
    ]

    for path in candidate_paths:
        value = read_reg_value(winreg.HKEY_CURRENT_USER, path, "CurrentVirtualDesktop")
        if value is None:
            continue

        if isinstance(value, (bytes, bytearray)) and len(value) == 16:
            g = guid_from_bytes_le(value)
            if g:
                return g

        if isinstance(value, str):
            g = parse_guidish_string(value)
            if g:
                return g

    return None

def read_desktop_id_list():
    """Read the ordered list of all known virtual desktop GUIDs."""
    path = r"Software\Microsoft\Windows\CurrentVersion\Explorer\VirtualDesktops"
    value = read_reg_value(winreg.HKEY_CURRENT_USER, path, "VirtualDesktopIDs")
    return parse_guid_list(value)

def get_current_desktop_number():
    """Map the current desktop GUID to a 1-based desktop index."""
    desktops = read_desktop_id_list()
    current = read_current_desktop_guid()

    log("desktops:", desktops)
    log("current :", current)

    if not desktops and not current:
        return 1
    if current and not desktops:
        return 1
    if desktops and not current:
        return None
    if current is not None and current in desktops:
        return desktops.index(current) + 1
    return None


# =========================
# Entry point
# =========================
if __name__ == "__main__":
    Overlay().run()