state_lock_overlay.py

import sys
import os
import json
import ctypes
from dataclasses import dataclass

from winreg import (
    ConnectRegistry,
    HKEY_CURRENT_USER,
    OpenKey,
    QueryValueEx,
)

from PyQt6.QtWidgets import QApplication, QWidget
from PyQt6.QtCore import Qt, QTimer, QPoint, QRectF
from PyQt6.QtGui import (
    QPainter,
    QColor,
    QBrush,
    QPen,
    QFont,
    QMouseEvent,
)

# ---------------------------------------------------------------------
# Windows virtual-key codes
# ---------------------------------------------------------------------
VK_CAPITAL = 0x14
VK_NUMLOCK = 0x90
VK_SCROLL = 0x91

# ---------------------------------------------------------------------
# Timers (ms)
# ---------------------------------------------------------------------
KEY_TIMER_MS = 150
THEME_TIMER_MS = 5000
TOP_TIMER_MS = 50
HIDE_TIMER_MS = 3000


@dataclass(frozen=True)
class Indicator:
    label: str
    vk_code: int
    x: int


class StateLockOverlay(QWidget):
    """
    Floating overlay widget that displays the state of
    Caps Lock, Num Lock and Scroll Lock.
    """

    CONFIG_FILENAME = ".state_lock_settings.json"

    DEFAULT_REL_X = 0.759
    DEFAULT_REL_Y = 0.962

    WIDTH = 130
    HEIGHT = 40

    def __init__(self):
        super().__init__()

        self.config_path = os.path.join(
            os.path.expanduser("~"),
            "Documents",
            self.CONFIG_FILENAME,
        )

        self._setup_window()
        self._setup_state()
        self._load_settings()
        self._setup_visuals()
        self._setup_screen()
        self._setup_timers()

        self.last_states = self.get_lock_states() 
        self.apply_theme()

    # ------------------------------------------------------------------
    # Setup
    # ------------------------------------------------------------------

    def _setup_window(self):
        self.setWindowFlags(
            Qt.WindowType.FramelessWindowHint
            | Qt.WindowType.WindowStaysOnTopHint
            | Qt.WindowType.Tool
        )
        self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground)
        self.setWindowOpacity(0.9)
        self.resize(self.WIDTH, self.HEIGHT)

    def _setup_state(self):
        self.dragging = False
        self.offset_pos = QPoint()
        self.origin_pos = QPoint()

        self.rel_x = self.DEFAULT_REL_X
        self.rel_y = self.DEFAULT_REL_Y
        self.locked = False
        self.hide_timeout_ms = HIDE_TIMER_MS

    def _setup_visuals(self):
        self.accent_color = QColor("#0078D7")
        self.bg_color = QColor("#F3F3F3")
        self.text_off = QColor("#999999")

        self.indicators = (
            Indicator("C", VK_CAPITAL, 10),
            Indicator("N", VK_NUMLOCK, 50),
            Indicator("S", VK_SCROLL, 90),
        )

        self.update_cursor()

    def _setup_screen(self):
        self.primary_screen = QApplication.primaryScreen()
        self.primary_screen.geometryChanged.connect(
            self._on_resolution_change
        )
        self._update_position_from_relative()

    def _setup_timers(self):
        self.key_timer = QTimer(self)
        self.key_timer.timeout.connect(self._poll_keys)
        self.key_timer.start(KEY_TIMER_MS)

        self.theme_timer = QTimer(self)
        self.theme_timer.timeout.connect(self.apply_theme)
        self.theme_timer.start(THEME_TIMER_MS)

        self.top_timer = QTimer(self)
        self.top_timer.timeout.connect(self.raise_)
        self.top_timer.start(TOP_TIMER_MS)

        self.hide_timer = QTimer(self)
        self.hide_timer.setSingleShot(True)
        self.hide_timer.timeout.connect(self._auto_hide)
        if self.locked: self.hide_timer.start(self.hide_timeout_ms)


    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def raise_(self):
        if self.isVisible():
            super().raise_() 
        else:
            return

    def _load_settings(self):
        if not os.path.exists(self.config_path):
            return

        try:
            with open(self.config_path, "r", encoding="utf-8") as f:
                data = json.load(f)

            self.rel_x = data.get("rel_x", self.DEFAULT_REL_X)
            self.rel_y = data.get("rel_y", self.DEFAULT_REL_Y)
            self.locked = data.get("locked", False)
            self.hide_timeout_ms = data.get(
                "hide_timer_in_ms", HIDE_TIMER_MS
            )
        except Exception as exc:
            print(f"Failed to load settings: {exc}")

    def _save_settings(self):
        data = {
            "rel_x": self.rel_x,
            "rel_y": self.rel_y,
            "locked": self.locked,
            "hide_timer_in_ms": self.hide_timeout_ms,
        }

        try:
            with open(self.config_path, "w", encoding="utf-8") as f:
                json.dump(data, f)
        except Exception as exc:
            print(f"Failed to save settings: {exc}")

    # ------------------------------------------------------------------
    # Positioning
    # ------------------------------------------------------------------

    def _on_resolution_change(self, _):
        self._update_position_from_relative()

    def _update_position_from_relative(self):
        geo = self.primary_screen.geometry()

        x = int(geo.width() * self.rel_x)
        y = int(geo.height() * self.rel_y)

        self.move(
            max(0, min(x, geo.width() - self.WIDTH)),
            max(0, min(y, geo.height() - self.HEIGHT)),
        )

    # ------------------------------------------------------------------
    # Painting
    # ------------------------------------------------------------------

    def paintEvent(self, _event):
        if not self.isVisible():
            return
        
        painter = QPainter(self)
        painter.setRenderHint(QPainter.RenderHint.Antialiasing)
        painter.setFont(QFont("Segoe UI Variable", 10, QFont.Weight.Bold))

        for indicator in self.indicators:
            active = self.is_key_active(indicator.vk_code)

            fill = self.accent_color if active else self.bg_color
            text = QColor("white") if active else self.text_off

            rect = QRectF(indicator.x, 8, 28, 24)

            painter.setPen(QPen(fill if active else self.text_off, 2))
            painter.setBrush(QBrush(fill))
            painter.drawEllipse(rect)

            painter.setPen(text)
            painter.drawText(rect, Qt.AlignmentFlag.AlignCenter, indicator.label)

    # ------------------------------------------------------------------
    # Mouse events
    # ------------------------------------------------------------------

    def mousePressEvent(self, event: QMouseEvent):
        if event.button() == Qt.MouseButton.LeftButton and not self.locked:
            self.dragging = True
            self.offset_pos = event.position().toPoint()
            self.origin_pos = self.pos()

        elif event.button() == Qt.MouseButton.MiddleButton:
            self.toggle_lock()

        elif event.button() == Qt.MouseButton.RightButton:
            self._save_settings()
            QApplication.quit()

    def mouseReleaseEvent(self, _event):
        if not self.dragging:
            return

        self.dragging = False
        geo = self.primary_screen.geometry()

        self.rel_x = self.x() / geo.width()
        self.rel_y = self.y() / geo.height()

        self._save_settings()

    def mouseMoveEvent(self, event: QMouseEvent):
        if not self.dragging or self.locked:
            return

        target = event.globalPosition().toPoint() - self.offset_pos

        if QApplication.keyboardModifiers() & Qt.KeyboardModifier.ShiftModifier:
            if abs(target.x() - self.origin_pos.x()) > abs(
                target.y() - self.origin_pos.y()
            ):
                target.setY(self.origin_pos.y())
            else:
                target.setX(self.origin_pos.x())

        geo = self.primary_screen.geometry()

        self.move(
            max(0, min(target.x(), geo.width() - self.WIDTH)),
            max(0, min(target.y(), geo.height() - self.HEIGHT)),
        )

    # ------------------------------------------------------------------
    # Lock & visibility
    # ------------------------------------------------------------------

    def toggle_lock(self):
        self.locked = not self.locked
        self.update_cursor()
        self._save_settings()

        if self.locked:
            self.hide_timer.start(self.hide_timeout_ms)

    def update_cursor(self):
        self.setCursor(
            Qt.CursorShape.ArrowCursor
            if self.locked
            else Qt.CursorShape.SizeAllCursor
        )

    def _auto_hide(self):
        if self.locked:
            self.hide()

    # ------------------------------------------------------------------
    # Windows integration
    # ------------------------------------------------------------------

    @staticmethod
    def _get_key_state(vk_code: int) -> int:
        return ctypes.windll.user32.GetKeyState(vk_code)

    def is_key_active(self, vk_code: int) -> bool:
        return bool(self._get_key_state(vk_code) & 1)

    def get_lock_states(self) -> dict:
        return {
            VK_CAPITAL: self.is_key_active(VK_CAPITAL),
            VK_NUMLOCK: self.is_key_active(VK_NUMLOCK),
            VK_SCROLL: self.is_key_active(VK_SCROLL),
        }

    def _poll_keys(self):
        current = self.get_lock_states()

        if current == self.last_states:
            return

        self.last_states = current

        if not self.isVisible():
            self.show()

        self.update()
        self.raise_()
        self.hide_timer.start(self.hide_timeout_ms)

    # ------------------------------------------------------------------
    # Theme
    # ------------------------------------------------------------------

    def get_windows_theme(self):
        try:
            reg = ConnectRegistry(None, HKEY_CURRENT_USER)

            dwm = OpenKey(reg, r"Software\Microsoft\Windows\DWM")
            color = QueryValueEx(dwm, "ColorizationColor")[0]
            accent = f"#{hex(color)[4:10]}"

            personalize = OpenKey(
                reg,
                r"Software\Microsoft\Windows\CurrentVersion\Themes\Personalize",
            )
            is_dark = QueryValueEx(personalize, "AppsUseLightTheme")[0] == 0

            return is_dark, accent if len(accent) == 7 else "#0078D7"

        except Exception:
            return True, "#0078D7"

    def apply_theme(self):
        is_dark, accent = self.get_windows_theme()

        self.accent_color = QColor(accent)
        self.bg_color = QColor("#1A1A1A") if is_dark else QColor("#F3F3F3")
        self.text_off = QColor("#666666") if is_dark else QColor("#999999")

        self.update()


# ---------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------
if __name__ == "__main__":
    app = QApplication(sys.argv)
    overlay = StateLockOverlay()
    overlay.show()
    sys.exit(app.exec())