soniox-mac-voice-typing/src/HotkeyManager.swift

import Cocoa
import Carbon
import ApplicationServices

public enum DictationMode: String, CaseIterable {
    case pushToTalk = "pushToTalk"  // Hold to record, release to transcribe
    case toggle = "toggle"          // Press once to start, press again to stop
    
    public var localizedTitle: String {
        switch self {
        case .pushToTalk:
            return "Push to Talk (Hold Key)"
        case .toggle:
            return "Toggle (Press to Start / Stop)"
        }
    }
}

public enum HotkeyPreset: String, CaseIterable {
    case micKey = "Microphone Key (🎙️ / F5)"
    case option = "Option (Hold ⌥)"
    case doubleTapFn = "Double-Tap Fn (Globe 🌐)"
    case holdFn = "Hold Fn (Globe 🌐)"
    case controlSpace = "Control + Space"
    case capsLock = "Caps Lock"
    case optionSpace = "Option + Space"
    case cmdShiftSpace = "Cmd + Shift + Space"
    case f8 = "F8"
    
    public var keyCode: UInt32 {
        switch self {
        case .micKey:
            return UInt32(kVK_F5) // 96
        case .controlSpace, .optionSpace, .cmdShiftSpace:
            return UInt32(kVK_Space) // 49
        case .option:
            return UInt32(kVK_Option) // 58
        case .doubleTapFn, .holdFn:
            return UInt32(kVK_Function) // 63
        case .capsLock:
            return UInt32(kVK_CapsLock) // 57
        case .f8:
            return UInt32(kVK_F8)
        }
    }
    
    public var carbonModifiers: UInt32 {
        switch self {
        case .controlSpace:
            return UInt32(controlKey)
        case .optionSpace:
            return UInt32(optionKey)
        case .cmdShiftSpace:
            return UInt32(cmdKey | shiftKey)
        case .micKey, .option, .doubleTapFn, .holdFn, .capsLock, .f8:
            return 0
        }
    }
}

public final class HotkeyManager {
    public static let shared = HotkeyManager()
    
    public var onHotkeyPressed: (() -> Void)?
    public var onHotkeyReleased: (() -> Void)?
    public var onEscapePressed: (() -> Void)? // Non-exclusive: both Soniox and active app (Telegram, etc.) receive Escape!
    
    private let stateLock = NSLock()
    private var hotKeyRef: EventHotKeyRef?
    private var eventHandlerRef: EventHandlerRef?
    private var eventTapPort: CFMachPort?
    private var runLoopSource: CFRunLoopSource?
    private var globalMonitor: Any?
    private var localMonitor: Any?
    private var escapePollTimer: Timer?
    private var wasEscapePhysicallyDown = false
    
    private var isOptionPhysicallyDown = false
    private var isCapsLockPhysicallyDown = false
    private var isF5PhysicallyDown = false
    private var isFnPhysicallyDown = false
    
    public var isPhysicalKeyDown: Bool {
        stateLock.lock()
        defer { stateLock.unlock() }
        return isOptionPhysicallyDown || isCapsLockPhysicallyDown || isF5PhysicallyDown || isFnPhysicallyDown
    }
    private var lastFnPressTime: TimeInterval = 0
    private var lastTriggerTime: TimeInterval = 0
    private var lastEscapeTime: TimeInterval = 0
    
    public var currentPreset: HotkeyPreset {
        get {
            let val = UserDefaults.standard.string(forKey: "SonioxHotkeyPreset") ?? HotkeyPreset.option.rawValue
            return HotkeyPreset(rawValue: val) ?? .option
        }
        set {
            UserDefaults.standard.set(newValue.rawValue, forKey: "SonioxHotkeyPreset")
            registerHotkeys()
        }
    }
    
    public var currentMode: DictationMode {
        get {
            let val = UserDefaults.standard.string(forKey: "SonioxDictationMode") ?? DictationMode.pushToTalk.rawValue
            return DictationMode(rawValue: val) ?? .pushToTalk
        }
        set {
            UserDefaults.standard.set(newValue.rawValue, forKey: "SonioxDictationMode")
            registerHotkeys()
        }
    }
    
    private init() {
        setupSleepWakeObservers()
    }
    
    private func setupSleepWakeObservers() {
        NSWorkspace.shared.notificationCenter.addObserver(forName: NSWorkspace.willSleepNotification, object: nil, queue: .main) { [weak self] _ in
            guard let self = self else { return }
            print("HotkeyManager: Mac entering sleep. Resetting keys & unregistering taps.")
            self.resetKeyState()
            self.unregisterHotkeys()
        }
        
        NSWorkspace.shared.notificationCenter.addObserver(forName: NSWorkspace.didWakeNotification, object: nil, queue: .main) { [weak self] _ in
            guard let self = self else { return }
            print("HotkeyManager: Mac woke from sleep. Re-registering hotkeys in 0.5s...")
            self.resetKeyState()
            DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) {
                self.registerHotkeys()
            }
        }
        
        NSWorkspace.shared.notificationCenter.addObserver(forName: NSWorkspace.screensDidWakeNotification, object: nil, queue: .main) { [weak self] _ in
            guard let self = self else { return }
            print("HotkeyManager: Screens woke. Verifying hotkey event tap...")
            if let tap = self.eventTapPort {
                CGEvent.tapEnable(tap: tap, enable: true)
            } else {
                self.registerHotkeys()
            }
        }
    }
    
    public func resetKeyState() {
        stateLock.lock()
        isOptionPhysicallyDown = false
        isF5PhysicallyDown = false
        isCapsLockPhysicallyDown = false
        isFnPhysicallyDown = false
        wasEscapePhysicallyDown = false
        stateLock.unlock()
    }
    
    /// Enables/disables ultra-lightweight 60Hz hardware Escape key state polling while Soniox is active (recording/transcribing/retrying/HUD visible).
    /// Because it reads hardware HID state directly (CGEventSource.keyState), it NEVER conflicts with or blocks Telegram or other apps' Escape hotkeys,
    /// and works 100% reliably regardless of TCC Input Monitoring ad-hoc signature state!
    public func setEscapeMonitoringActive(_ active: Bool) {
        DispatchQueue.main.async { [weak self] in
            guard let self = self else { return }
            self.escapePollTimer?.invalidate()
            self.escapePollTimer = nil
            self.wasEscapePhysicallyDown = false
            
            if active {
                self.escapePollTimer = Timer.scheduledTimer(withTimeInterval: 0.016, repeats: true) { [weak self] _ in
                    guard let self = self else { return }
                    let isDown = CGEventSource.keyState(.combinedSessionState, key: 53) || CGEventSource.keyState(.hidSystemState, key: 53)
                    if isDown && !self.wasEscapePhysicallyDown {
                        self.wasEscapePhysicallyDown = true
                        self.triggerEscape()
                    } else if !isDown {
                        self.wasEscapePhysicallyDown = false
                    }
                }
            }
        }
    }
    
    private func triggerEscape() {
        stateLock.lock()
        let now = Date().timeIntervalSince1970
        guard (now - lastEscapeTime) > 0.15 else {
            stateLock.unlock()
            return
        }
        lastEscapeTime = now
        stateLock.unlock()
        
        DispatchQueue.main.async { [weak self] in
            self?.onEscapePressed?()
        }
    }
    
    public func registerHotkeys() {
        unregisterHotkeys()
        resetKeyState()
        
        let preset = currentPreset
        print("HotkeyManager: Registering hotkey for '\(preset.rawValue)', mode: '\(currentMode.rawValue)'")
        
        // 1. Carbon HotKey for multi-key combos (Control+Space, Option+Space, Cmd+Shift+Space, F8)
        if preset != .option && preset != .capsLock && preset != .micKey && preset != .doubleTapFn && preset != .holdFn {
            var eventTypes = [
                EventTypeSpec(eventClass: OSType(kEventClassKeyboard), eventKind: UInt32(kEventHotKeyPressed)),
                EventTypeSpec(eventClass: OSType(kEventClassKeyboard), eventKind: UInt32(kEventHotKeyReleased))
            ]
            
            let selfPtr = Unmanaged.passUnretained(self).toOpaque()
            let handlerCallback: EventHandlerUPP = { (_, eventRef, userData) -> OSStatus in
                guard let eventRef = eventRef, let userData = userData else { return noErr }
                let manager = Unmanaged<HotkeyManager>.fromOpaque(userData).takeUnretainedValue()
                
                let kind = GetEventKind(eventRef)
                if kind == UInt32(kEventHotKeyPressed) {
                    manager.triggerAction()
                } else if kind == UInt32(kEventHotKeyReleased) {
                    if manager.currentMode == .pushToTalk {
                        DispatchQueue.main.async {
                            manager.onHotkeyReleased?()
                        }
                    }
                }
                return noErr
            }
            
            InstallEventHandler(GetApplicationEventTarget(), handlerCallback, 2, &eventTypes, selfPtr, &eventHandlerRef)
            
            let hotKeyID = EventHotKeyID(signature: OSType(0x534F4E58), id: 1)
            RegisterEventHotKey(
                preset.keyCode,
                preset.carbonModifiers,
                hotKeyID,
                GetApplicationEventTarget(),
                0,
                &hotKeyRef
            )
        }
        
        // 2. Global Event Tap for high-priority interception
        setupEventTap()
        
        // 3. Always-active non-exclusive NSEvent monitor for modifier & Escape pass-through
        setupGlobalMonitor()
    }
    
    private func triggerAction() {
        stateLock.lock()
        let now = Date().timeIntervalSince1970
        guard (now - lastTriggerTime) > 0.10 else {
            stateLock.unlock()
            return
        }
        lastTriggerTime = now
        stateLock.unlock()
        
        DispatchQueue.main.async { [weak self] in
            self?.onHotkeyPressed?()
        }
    }
    
    private func setupEventTap() {
        let mask = (1 << CGEventType.flagsChanged.rawValue) | (1 << CGEventType.keyDown.rawValue) | (1 << CGEventType.keyUp.rawValue)
        let selfPtr = Unmanaged.passUnretained(self).toOpaque()
        
        let callback: CGEventTapCallBack = { (proxy, type, event, refcon) -> Unmanaged<CGEvent>? in
            guard let refcon = refcon else { return Unmanaged.passRetained(event) }
            let manager = Unmanaged<HotkeyManager>.fromOpaque(refcon).takeUnretainedValue()
            
            // Auto-recover if macOS disabled tap on timeout or user input
            if type == .tapDisabledByTimeout || type == .tapDisabledByUserInput {
                print("HotkeyManager: EventTap was disabled by OS (\(type.rawValue)). Re-enabling now...")
                if let tap = manager.eventTapPort {
                    CGEvent.tapEnable(tap: tap, enable: true)
                }
                return nil
            }
            
            let flags = event.flags.rawValue
            let keyCode = event.getIntegerValueField(.keyboardEventKeycode)
            
            // 0. Escape Key (kVK_Escape = 53): Non-blocking pass-through!
            // Both Soniox Voice AND the active app (Telegram, Chrome, etc.) execute their Escape action simultaneously!
            if keyCode == 53 && type == .keyDown {
                manager.triggerEscape()
                return Unmanaged.passRetained(event) // NEVER swallow Escape so Telegram/apps also receive it!
            }
            
            // 1. Microphone Key (🎙️ / F5)
            if manager.currentPreset == .micKey {
                if keyCode == 96 { // 96 = kVK_F5
                    if type == .keyDown {
                        manager.stateLock.lock()
                        let wasDown = manager.isF5PhysicallyDown
                        manager.isF5PhysicallyDown = true
                        manager.stateLock.unlock()
                        if !wasDown {
                            manager.triggerAction()
                        }
                        return nil // Swallow F5 so it doesn't leak to apps
                    } else if type == .keyUp {
                        manager.stateLock.lock()
                        manager.isF5PhysicallyDown = false
                        manager.stateLock.unlock()
                        if manager.currentMode == .pushToTalk {
                            DispatchQueue.main.async {
                                manager.onHotkeyReleased?()
                            }
                        }
                        return nil // Swallow F5 keyUp
                    }
                }
            }
            
            // 2. Option Key
            else if manager.currentPreset == .option {
                if keyCode == 58 || keyCode == 61 { // 58 = Left Option, 61 = Right Option
                    let isAlt = (flags & CGEventFlags.maskAlternate.rawValue) != 0
                    manager.stateLock.lock()
                    let changed = (isAlt != manager.isOptionPhysicallyDown)
                    if changed {
                        manager.isOptionPhysicallyDown = isAlt
                    }
                    manager.stateLock.unlock()
                    
                    if changed {
                        if isAlt {
                            manager.triggerAction()
                        } else if manager.currentMode == .pushToTalk {
                            DispatchQueue.main.async {
                                manager.onHotkeyReleased?()
                            }
                        }
                    }
                }
            }
            
            // 3. CapsLock Key
            else if manager.currentPreset == .capsLock {
                if keyCode == 57 { // 57 = CapsLock
                    manager.stateLock.lock()
                    let nowDown = !manager.isCapsLockPhysicallyDown
                    manager.isCapsLockPhysicallyDown = nowDown
                    manager.stateLock.unlock()
                    
                    if nowDown {
                        manager.triggerAction()
                    } else {
                        if manager.currentMode == .pushToTalk {
                            DispatchQueue.main.async {
                                manager.onHotkeyReleased?()
                            }
                        }
                    }
                    return nil
                }
            }
            
            // 4. Control + Space
            else if manager.currentPreset == .controlSpace {
                if keyCode == 49 && type == .keyDown { // 49 = Space
                    let isCtrl = (flags & CGEventFlags.maskControl.rawValue) != 0
                    if isCtrl {
                        manager.triggerAction()
                        return nil
                    }
                }
            }
            
            // 5. Fn / Globe Key (Double-Tap or Hold)
            else if manager.currentPreset == .doubleTapFn || manager.currentPreset == .holdFn {
                if keyCode == 63 { // 63 = kVK_Function (Fn / Globe)
                    let isFnActive = (flags & CGEventFlags.maskSecondaryFn.rawValue) != 0
                    manager.stateLock.lock()
                    let changed = (isFnActive != manager.isFnPhysicallyDown)
                    if changed {
                        manager.isFnPhysicallyDown = isFnActive
                    }
                    manager.stateLock.unlock()
                    
                    if changed {
                        if isFnActive {
                            let now = Date().timeIntervalSince1970
                            let interval = now - manager.lastFnPressTime
                            manager.lastFnPressTime = now
                            
                            if manager.currentPreset == .doubleTapFn {
                                if interval > 0.05 && interval < 0.45 {
                                    manager.lastFnPressTime = 0
                                    manager.triggerAction()
                                    return nil
                                }
                            } else if manager.currentPreset == .holdFn {
                                manager.triggerAction()
                                return nil
                            }
                        } else {
                            if manager.currentPreset == .holdFn && manager.currentMode == .pushToTalk {
                                DispatchQueue.main.async {
                                    manager.onHotkeyReleased?()
                                }
                                return nil
                            }
                        }
                    }
                }
            }
            
            return Unmanaged.passRetained(event)
        }
        
        if let tap = CGEvent.tapCreate(
            tap: .cghidEventTap,
            place: .headInsertEventTap,
            options: .defaultTap,
            eventsOfInterest: CGEventMask(mask),
            callback: callback,
            userInfo: selfPtr
        ) {
            self.eventTapPort = tap
            let source = CFMachPortCreateRunLoopSource(kCFAllocatorDefault, tap, 0)
            self.runLoopSource = source
            CFRunLoopAddSource(CFRunLoopGetMain(), source, .commonModes)
            CGEvent.tapEnable(tap: tap, enable: true)
            print("HotkeyManager: CGEventTap registered successfully.")
        }
    }
    
    private func setupGlobalMonitor() {
        if let mon = globalMonitor {
            NSEvent.removeMonitor(mon)
            self.globalMonitor = nil
        }
        if let lmon = localMonitor {
            NSEvent.removeMonitor(lmon)
            self.localMonitor = nil
        }
        
        // Local monitor if Soniox menu/HUD is ever key
        localMonitor = NSEvent.addLocalMonitorForEvents(matching: [.keyDown, .flagsChanged]) { [weak self] event in
            if event.keyCode == 53 && event.type == .keyDown {
                self?.triggerEscape()
            }
            return event // Pass through
        }
        
        // Global passive monitor: always active alongside EventTap (deduplicated via stateLock)
        globalMonitor = NSEvent.addGlobalMonitorForEvents(matching: [.flagsChanged, .keyDown, .keyUp]) { [weak self] event in
            guard let self = self else { return }
            
            // Wake up eventTap if it was disabled by OS
            if let tap = self.eventTapPort, !CGEvent.tapIsEnabled(tap: tap) {
                CGEvent.tapEnable(tap: tap, enable: true)
            }
            
            let flags = event.modifierFlags
            let keyCode = event.keyCode
            
            // Non-exclusive Escape pass-through
            if keyCode == 53 && event.type == .keyDown {
                self.triggerEscape()
                return
            }
            
            if self.currentPreset == .option {
                if keyCode == 58 || keyCode == 61 {
                    let isAlt = flags.contains(.option)
                    self.stateLock.lock()
                    let changed = (isAlt != self.isOptionPhysicallyDown)
                    if changed { self.isOptionPhysicallyDown = isAlt }
                    self.stateLock.unlock()
                    
                    if changed {
                        if isAlt {
                            self.triggerAction()
                        } else if self.currentMode == .pushToTalk {
                            DispatchQueue.main.async {
                                self.onHotkeyReleased?()
                            }
                        }
                    }
                }
            } else if self.currentPreset == .controlSpace {
                if keyCode == 49 && flags.contains(.control) {
                    self.triggerAction()
                }
            } else if self.currentPreset == .capsLock {
                if keyCode == 57 {
                    let isCaps = flags.contains(.capsLock)
                    self.stateLock.lock()
                    let changed = (isCaps != self.isCapsLockPhysicallyDown)
                    if changed { self.isCapsLockPhysicallyDown = isCaps }
                    self.stateLock.unlock()
                    
                    if changed {
                        if isCaps {
                            self.triggerAction()
                        } else if self.currentMode == .pushToTalk {
                            DispatchQueue.main.async {
                                self.onHotkeyReleased?()
                            }
                        }
                    }
                }
            }
        }
    }
    
    public func unregisterHotkeys() {
        resetKeyState()
        escapePollTimer?.invalidate()
        escapePollTimer = nil
        if let ref = hotKeyRef {
            UnregisterEventHotKey(ref)
            hotKeyRef = nil
        }
        if let handler = eventHandlerRef {
            RemoveEventHandler(handler)
            eventHandlerRef = nil
        }
        if let tap = eventTapPort, let source = runLoopSource {
            CGEvent.tapEnable(tap: tap, enable: false)
            CFRunLoopRemoveSource(CFRunLoopGetMain(), source, .commonModes)
            self.eventTapPort = nil
            self.runLoopSource = nil
        }
        if let mon = globalMonitor {
            NSEvent.removeMonitor(mon)
            self.globalMonitor = nil
        }
        if let lmon = localMonitor {
            NSEvent.removeMonitor(lmon)
            self.localMonitor = nil
        }
    }
}