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

import Cocoa
import ApplicationServices

public final class InputBoxManager {
    public static let shared = InputBoxManager()
    
    private var targetedElement: AXUIElement?
    private var streamedCharacterCount: Int = 0
    private var didStreamAnyText: Bool = false
    private let lock = NSLock()
    
    private init() {}
    
    /// Prepares and guarantees that an input box is focused before or at the start of recording
    @discardableResult
    public func ensureFocusedInputBox() -> AXUIElement? {
        lock.lock()
        defer { lock.unlock() }
        
        streamedCharacterCount = 0
        didStreamAnyText = false
        targetedElement = nil
        
        guard AXIsProcessTrusted() else {
            print("InputBoxManager: AX not trusted, falling back to general keyboard events")
            return nil
        }
        
        guard let frontApp = NSWorkspace.shared.frontmostApplication else { return nil }
        let appElem = AXUIElementCreateApplication(frontApp.processIdentifier)
        
        // 1. Check if an element is already focused in the frontmost application
        var focusedObj: CFTypeRef?
        if AXUIElementCopyAttributeValue(appElem, kAXFocusedUIElementAttribute as CFString, &focusedObj) == .success,
           let focused = focusedObj {
            let elem = (focused as! AXUIElement)
            if isEditableElement(elem) {
                targetedElement = elem
                return elem
            }
        }
        
        // 2. Also check systemWide focus
        let sysWide = AXUIElementCreateSystemWide()
        if AXUIElementCopyAttributeValue(sysWide, kAXFocusedUIElementAttribute as CFString, &focusedObj) == .success,
           let focused = focusedObj {
            let elem = (focused as! AXUIElement)
            if isEditableElement(elem) {
                targetedElement = elem
                return elem
            }
        }
        
        // 3. If no input is focused, search the active window of the frontmost app for an input field
        var windowObj: CFTypeRef?
        if AXUIElementCopyAttributeValue(appElem, kAXFocusedWindowAttribute as CFString, &windowObj) == .success,
           let win = windowObj {
            let winElem = (win as! AXUIElement)
            if let bestInput = findBestInputElement(in: winElem) {
                // Focus it via Accessibility
                AXUIElementSetAttributeValue(bestInput, kAXFocusedAttribute as CFString, kCFBooleanTrue)
                
                // If it has a clickable position, click its center to ensure real input focus
                if let center = getElementCenter(bestInput) {
                    clickAtPoint(center)
                }
                
                targetedElement = bestInput
                return bestInput
            }
        }
        
        return nil
    }
    
    /// Stream delta text (incremental words) live into the focused input box
    public func streamDelta(text: String) {
        guard !text.isEmpty else { return }
        
        lock.lock()
        didStreamAnyText = true
        streamedCharacterCount += text.count
        let target = targetedElement
        lock.unlock()
        
        // 1. Try direct Accessibility insertion if element supports it
        if let elem = target {
            var selectedTextObj: CFTypeRef?
            if AXUIElementCopyAttributeValue(elem, kAXSelectedTextAttribute as CFString, &selectedTextObj) == .success {
                let err = AXUIElementSetAttributeValue(elem, kAXSelectedTextAttribute as CFString, text as CFString)
                if err == .success {
                    return
                }
            }
        }
        
        // 2. Native Unicode CGEvent typing (universal, works in Telegram, browsers, Electron, native apps)
        typeUnicodeString(text)
    }
    
    /// Commits final text to the input box and system clipboard
    public func commitFinalText(processedText: String, addTrailingSpace: Bool) {
        guard !processedText.isEmpty else { return }
        
        let textToSave = addTrailingSpace ? (processedText + " ") : processedText
        
        // 1. Always put full processed text into clipboard
        let pb = NSPasteboard.general
        pb.clearContents()
        pb.setString(textToSave, forType: .string)
        
        lock.lock()
        let hadStreamed = didStreamAnyText
        let streamedCount = streamedCharacterCount
        targetedElement = nil
        lock.unlock()
        
        // If we didn't stream live (e.g. command or burst retry), paste the full text now
        if !hadStreamed || streamedCount == 0 {
            pasteClipboardToFrontmostApp()
        }
    }
    
    /// Simulates Cmd+V to paste clipboard contents
    public func pasteClipboardToFrontmostApp() {
        DispatchQueue.main.asyncAfter(deadline: .now() + 0.02) {
            let source = CGEventSource(stateID: .combinedSessionState)
            let vKeyCode: CGKeyCode = 0x09 // 'v'
            if let down = CGEvent(keyboardEventSource: source, virtualKey: vKeyCode, keyDown: true),
               let up = CGEvent(keyboardEventSource: source, virtualKey: vKeyCode, keyDown: false) {
                down.flags = .maskCommand
                up.flags = .maskCommand
                down.post(tap: .cghidEventTap)
                up.post(tap: .cghidEventTap)
            }
        }
    }
    
    // MARK: - Helpers
    
    private func isEditableElement(_ elem: AXUIElement) -> Bool {
        var roleObj: CFTypeRef?
        AXUIElementCopyAttributeValue(elem, kAXRoleAttribute as CFString, &roleObj)
        let role = (roleObj as? String) ?? ""
        if role == "AXTextField" || role == "AXTextArea" || role == "AXComboBox" || role == "AXWebArea" {
            return true
        }
        var valObj: CFTypeRef?
        if AXUIElementCopyAttributeValue(elem, kAXValueAttribute as CFString, &valObj) == .success {
            return true
        }
        return false
    }
    
    private func findBestInputElement(in parent: AXUIElement, maxDepth: Int = 8, depth: Int = 0) -> AXUIElement? {
        if depth > maxDepth { return nil }
        
        var roleObj: CFTypeRef?
        AXUIElementCopyAttributeValue(parent, kAXRoleAttribute as CFString, &roleObj)
        let role = (roleObj as? String) ?? ""
        
        if role == "AXTextField" || role == "AXTextArea" || role == "AXComboBox" {
            return parent
        }
        
        var childrenObj: CFTypeRef?
        if AXUIElementCopyAttributeValue(parent, kAXChildrenAttribute as CFString, &childrenObj) == .success,
           let children = childrenObj as? [AXUIElement] {
            for child in children {
                if let found = findBestInputElement(in: child, maxDepth: maxDepth, depth: depth + 1) {
                    return found
                }
            }
        }
        return nil
    }
    
    private func getElementCenter(_ elem: AXUIElement) -> CGPoint? {
        var posVal: CFTypeRef?
        var sizeVal: CFTypeRef?
        guard AXUIElementCopyAttributeValue(elem, kAXPositionAttribute as CFString, &posVal) == .success,
              AXUIElementCopyAttributeValue(elem, kAXSizeAttribute as CFString, &sizeVal) == .success else {
            return nil
        }
        var pt = CGPoint.zero
        var sz = CGSize.zero
        guard AXValueGetValue(posVal as! AXValue, .cgPoint, &pt),
              AXValueGetValue(sizeVal as! AXValue, .cgSize, &sz) else {
            return nil
        }
        return CGPoint(x: pt.x + sz.width / 2.0, y: pt.y + sz.height / 2.0)
    }
    
    private func clickAtPoint(_ point: CGPoint) {
        let source = CGEventSource(stateID: .hidSystemState)
        if let clickDown = CGEvent(mouseEventSource: source, mouseType: .leftMouseDown, mouseCursorPosition: point, mouseButton: .left),
           let clickUp = CGEvent(mouseEventSource: source, mouseType: .leftMouseUp, mouseCursorPosition: point, mouseButton: .left) {
            clickDown.post(tap: .cghidEventTap)
            clickUp.post(tap: .cghidEventTap)
        }
    }
    
    private func typeUnicodeString(_ string: String) {
        let source = CGEventSource(stateID: .hidSystemState)
        let utf16 = Array(string.utf16)
        
        guard let keyDown = CGEvent(keyboardEventSource: source, virtualKey: 0, keyDown: true),
              let keyUp = CGEvent(keyboardEventSource: source, virtualKey: 0, keyDown: false) else {
            return
        }
        
        keyDown.keyboardSetUnicodeString(stringLength: utf16.count, unicodeString: utf16)
        keyUp.keyboardSetUnicodeString(stringLength: utf16.count, unicodeString: utf16)
        
        keyDown.post(tap: .cghidEventTap)
        keyUp.post(tap: .cghidEventTap)
    }
}