cmd-key-happy in Swift

cmd-key-happy.swift

import Cocoa
import CoreGraphics
import Foundation

class CmdKeyHappy {
    func cleanup() {
        print("Cleaning up CmdKeyHappy...")
        for (pid, tap) in tapMap {
            print("Removing tap for PID: \(pid)")
            CFMachPortInvalidate(tap)
            if let runLoopSource = CFMachPortCreateRunLoopSource(nil, tap, 0) {
                CFRunLoopRemoveSource(CFRunLoopGetMain(), runLoopSource, .commonModes)
            }
        }
        tapMap.removeAll()
        NSWorkspace.shared.notificationCenter.removeObserver(self)
    }

    private var tapMap: [pid_t: CFMachPort] = [:]
    private let knownApps: Set<String>

    // Static callback function that will be used for all taps.
    private static let eventCallback: CGEventTapCallBack = { proxy, type, event, userInfo in
        guard type == .keyDown else {
            return Unmanaged.passUnretained(event)
        }

        guard let frontmostApp = NSWorkspace.shared.frontmostApplication,
              let appName = frontmostApp.localizedName else {
            return Unmanaged.passUnretained(event)
        }

        let knownApps = unsafeBitCast(userInfo, to: Set<String>.self)
        guard knownApps.contains(appName) else {
            return Unmanaged.passUnretained(event)
        }

        let flags = event.flags
        let isCommandPressed = flags.contains(.maskCommand)
        let isOptionPressed = flags.contains(.maskAlternate)

        guard isCommandPressed || isOptionPressed else {
            return Unmanaged.passUnretained(event)
        }

        let isAnyOtherModifierPressed = flags.contains(.maskShift) || flags.contains(.maskControl)

        // Early return if only Command is pressed
        if isCommandPressed && !isOptionPressed && !isAnyOtherModifierPressed && event.getIntegerValueField(.keyboardEventKeycode) == 0 {
            print("Command pressed alone; skipping.")
            return Unmanaged.passUnretained(event)
        }

        // Early return if both Command and Option are pressed
        if isCommandPressed && isOptionPressed {
            print("Both Command and Option pressed; skipping.")
            return Unmanaged.passUnretained(event)
        }


        var newFlags = flags
        if flags.contains(.maskCommand) {
            newFlags.remove(.maskCommand)
            newFlags.insert(.maskAlternate)
        } else if flags.contains(.maskAlternate) {
            newFlags.remove(.maskAlternate)
            newFlags.insert(.maskCommand)
        }

        event.flags = newFlags
        return Unmanaged.passUnretained(event)
    }

    init(knownApps: [String]) {
        self.knownApps = Set(knownApps)

        // First scan all running applications and create taps for known apps
        for app in NSWorkspace.shared.runningApplications {
            if let appName = app.localizedName, knownApps.contains(appName) {
                print("Found running known app '\(appName)'. Creating event tap.")
                tapApp(for: app.processIdentifier)
            }
        }

        NSWorkspace.shared.notificationCenter.addObserver(
          self,
          selector: #selector(handleAppActivated(_:)),
          name: NSWorkspace.didActivateApplicationNotification,
          object: nil
        )

        NSWorkspace.shared.notificationCenter.addObserver(
          self,
          selector: #selector(handleAppTerminated(_:)),
          name: NSWorkspace.didTerminateApplicationNotification,
          object: nil
        )
    }

    deinit {
        cleanup()
    }

    @objc private func handleAppActivated(_ notification: Notification) {
        guard let app = notification.userInfo?[NSWorkspace.applicationUserInfoKey] as? NSRunningApplication,
              let appName = app.localizedName else {
            return
        }

        let pid = app.processIdentifier

        if knownApps.contains(appName) {
            print("Known app '\(appName)' activated. Enabling its event tap.")
            if let tap = tapMap[pid] {
                CGEvent.tapEnable(tap: tap, enable: true)
            } else {
                tapApp(for: pid)
            }
        } else {
            print("App '\(appName)' not registered for key swapping.") // More accurate message
            disableAllEventTaps()
        }
    }

    @objc private func handleAppTerminated(_ notification: Notification) {
        guard let app = notification.userInfo?[NSWorkspace.applicationUserInfoKey] as? NSRunningApplication else {
            return
        }

        let pid = app.processIdentifier
        if let tap = tapMap[pid] {
            CFMachPortInvalidate(tap)
            tapMap.removeValue(forKey: pid)
            print("Event tap removed for terminated app.")
        }
    }

    private func disableAllEventTaps() {
        for (_, tap) in tapMap {
            CGEvent.tapEnable(tap: tap, enable: false)
        }
    }

    private func tapApp(for pid: pid_t) {
        if tapMap[pid] != nil {
            return
        }

        let eventMask: CGEventMask = (1 << CGEventType.keyDown.rawValue | 1 << CGEventType.flagsChanged.rawValue)

        // Pass knownApps set through userInfo
        let userInfo = unsafeBitCast(knownApps, to: UnsafeMutableRawPointer.self)

        guard let tap = CGEvent.tapCreate(
                tap: .cgAnnotatedSessionEventTap,
                place: .headInsertEventTap,
                options: .defaultTap,
                eventsOfInterest: eventMask,
                callback: CmdKeyHappy.eventCallback,
                userInfo: userInfo
              ) else {
            print("Failed to create event tap")
            return
        }

        let runLoopSource = CFMachPortCreateRunLoopSource(nil, tap, 0)
        CFRunLoopAddSource(CFRunLoopGetMain(), runLoopSource, .commonModes)

        tapMap[pid] = tap
        print("Event tap created for PID \(pid)")
    }
}

let cmdKeyHappy = CmdKeyHappy(knownApps: ["Ghostty", "Alacritty"])

signal(SIGTERM) { _ in
    print("Received SIGTERM, cleaning up...")
    DispatchQueue.main.async {
        cmdKeyHappy.cleanup()
        exit(0)
    }
}

signal(SIGINT) { _ in
    print("Received SIGINT, cleaning up...")
    DispatchQueue.main.async {
        cmdKeyHappy.cleanup()
        exit(0)
    }
}

RunLoop.main.run()