wisper-key/Sources/WhisperKey/NativeCompositionBridge.swift

import AppKit
import Carbon
import Foundation

@MainActor
final class NativeCompositionBridge: NSObject {
    private enum Message {
        static let ping = Notification.Name("com.romanv.WhisperKeyInput.ping")
        static let ready = Notification.Name("com.romanv.WhisperKeyInput.ready")
        static let draft = Notification.Name("com.romanv.WhisperKeyInput.draft")
        static let commit = Notification.Name("com.romanv.WhisperKeyInput.commit")
        static let cancel = Notification.Name("com.romanv.WhisperKeyInput.cancel")
        static let textKey = "text"
        static let activeKey = "active"
    }

    private static let inputSourceID = "com.romanv.inputmethod.WhisperKey"
    private static let installedBundle = FileManager.default.homeDirectoryForCurrentUser
        .appendingPathComponent("Library/Input Methods/WhisperKeyInput.app", isDirectory: true)

    var onAvailabilityChange: (() -> Void)?
    private var previousInputSourceID: String?
    private var restoreWorkItem: DispatchWorkItem?
    private(set) var inputMethodActive = false {
        didSet {
            if oldValue != inputMethodActive { onAvailabilityChange?() }
        }
    }

    var isAvailable: Bool {
        inputMethodActive
    }

    var isInstalled: Bool {
        FileManager.default.fileExists(atPath: Self.installedBundle.path)
    }

    func start() {
        DistributedNotificationCenter.default().addObserver(
            self,
            selector: #selector(handleReady),
            name: Message.ready,
            object: nil
        )
        if inputSource(withID: Self.inputSourceID) == nil {
            registerInstalledInputMethod()
        }
        launchInputMethod()
        ping()
    }

    func activateForDictation() -> Bool {
        restoreWorkItem?.cancel()
        restoreWorkItem = nil
        launchInputMethod()

        guard let source = inputSource(withID: Self.inputSourceID) else { return false }
        let currentID = currentInputSourceID()
        if currentID != Self.inputSourceID {
            previousInputSourceID = currentID
        }
        guard TISSelectInputSource(source) == noErr else {
            restoreInputSource()
            return false
        }
        wakeInputMethod()
        ping()
        return true
    }

    func ping() {
        DistributedNotificationCenter.default().post(
            name: Message.ping,
            object: nil,
            userInfo: nil
        )
    }

    @discardableResult
    func updateDraft(_ text: String) -> Bool {
        guard isAvailable else { return false }
        post(name: Message.draft, text: text)
        return true
    }

    @discardableResult
    func commit(_ text: String) -> Bool {
        let available = isAvailable
        if available { post(name: Message.commit, text: text) }
        scheduleInputSourceRestore()
        return available
    }

    func cancel(restoreInputSource: Bool = true) {
        if inputMethodActive {
            DistributedNotificationCenter.default().post(
                name: Message.cancel,
                object: nil,
                userInfo: nil
            )
        }
        if restoreInputSource { scheduleInputSourceRestore() }
    }

    func restoreInputSource() {
        restoreWorkItem?.cancel()
        restoreWorkItem = nil
        guard let previousInputSourceID else { return }
        self.previousInputSourceID = nil
        guard let source = inputSource(withID: previousInputSourceID) else { return }
        _ = TISSelectInputSource(source)
    }

    private func post(name: Notification.Name, text: String) {
        DistributedNotificationCenter.default().post(
            name: name,
            object: nil,
            userInfo: [Message.textKey: text]
        )
    }

    @objc private func handleReady(_ notification: Notification) {
        inputMethodActive = notification.userInfo?[Message.activeKey] as? Bool ?? false
    }

    private func registerInstalledInputMethod() {
        guard isInstalled else { return }
        _ = TISRegisterInputSource(Self.installedBundle as CFURL)
    }

    private func launchInputMethod() {
        guard isInstalled else { return }
        let configuration = NSWorkspace.OpenConfiguration()
        configuration.activates = false
        NSWorkspace.shared.openApplication(
            at: Self.installedBundle,
            configuration: configuration
        )
    }

    private func scheduleInputSourceRestore() {
        restoreWorkItem?.cancel()
        let work = DispatchWorkItem { [weak self] in
            self?.restoreInputSource()
        }
        restoreWorkItem = work
        DispatchQueue.main.asyncAfter(deadline: .now() + 0.3, execute: work)
    }

    private func wakeInputMethod() {
        guard let source = CGEventSource(stateID: .hidSystemState) else { return }
        let flags: CGEventFlags = [
            .maskCommand,
            .maskAlternate,
            .maskControl,
            .maskShift,
        ]
        guard let keyDown = CGEvent(
            keyboardEventSource: source,
            virtualKey: CGKeyCode(kVK_F20),
            keyDown: true
        ), let keyUp = CGEvent(
            keyboardEventSource: source,
            virtualKey: CGKeyCode(kVK_F20),
            keyDown: false
        ) else { return }
        keyDown.flags = flags
        keyUp.flags = flags
        keyDown.post(tap: .cghidEventTap)
        keyUp.post(tap: .cghidEventTap)
    }

    private func inputSource(withID identifier: String) -> TISInputSource? {
        guard let sources = TISCreateInputSourceList(nil, true)?.takeRetainedValue()
            as? [TISInputSource] else {
            return nil
        }
        return sources.first { inputSourceID(for: $0) == identifier }
    }

    private func currentInputSourceID() -> String? {
        guard let source = TISCopyCurrentKeyboardInputSource()?.takeRetainedValue() else { return nil }
        return inputSourceID(for: source)
    }

    private func inputSourceID(for source: TISInputSource) -> String? {
        guard let rawValue = TISGetInputSourceProperty(
            source,
            kTISPropertyInputSourceID
        ) else { return nil }
        return Unmanaged<CFString>
            .fromOpaque(rawValue)
            .takeUnretainedValue() as String
    }
}