air-accept/Sources/AirAccept/AccessibilityReader.swift

import AirAcceptCore
import AppKit
import ApplicationServices
import Foundation

struct PromptCandidateWithButton {
    var candidate: PromptCandidate
    var match: AirPlayPromptMatch
    var acceptButton: AXUIElement
    var acceptActionName: String
}

enum AccessibilityReader {
    private static let notificationCenterBundleID = "com.apple.notificationcenterui"

    static func isTrusted(prompt: Bool) -> Bool {
        let options = [kAXTrustedCheckOptionPrompt.takeUnretainedValue() as String: prompt] as CFDictionary
        return AXIsProcessTrustedWithOptions(options)
    }

    static func notificationCenterCandidate() -> PromptCandidateWithButton? {
        guard let window = notificationCenterWindow() else {
            return nil
        }

        return LocalTreeSearch.findCandidate(
            in: window,
            hostProcess: notificationCenterBundleID,
            maxDepth: 9,
            maxNodes: 450
        )
    }

    static func press(_ element: AXUIElement, actionName: String) -> AXError {
        AXUIElementPerformAction(element, actionName as CFString)
    }

    private static func notificationCenterWindow() -> AXUIElement? {
        guard
            let app = NSRunningApplication
                .runningApplications(withBundleIdentifier: notificationCenterBundleID)
                .first
        else {
            return nil
        }

        let appElement = AXUIElementCreateApplication(app.processIdentifier)
        var windowsRef: CFTypeRef?
        guard
            AXUIElementCopyAttributeValue(appElement, kAXWindowsAttribute as CFString, &windowsRef) == .success,
            let windows = windowsRef as? [AXUIElement],
            let window = windows.first
        else {
            return nil
        }

        return window
    }
}

private enum LocalTreeSearch {
    static func findCandidate(
        in root: AXUIElement,
        hostProcess: String,
        maxDepth: Int,
        maxNodes: Int
    ) -> PromptCandidateWithButton? {
        var visited = 0
        return visit(
            root,
            hostProcess: hostProcess,
            depth: 0,
            maxDepth: maxDepth,
            maxNodes: maxNodes,
            visited: &visited
        )
    }

    private static func visit(
        _ element: AXUIElement,
        hostProcess: String,
        depth: Int,
        maxDepth: Int,
        maxNodes: Int,
        visited: inout Int
    ) -> PromptCandidateWithButton? {
        guard depth <= maxDepth, visited < maxNodes else {
            return nil
        }
        visited += 1

        let snapshot = SnapshotBuilder.build(from: element, maxDepth: 4, maxNodes: 90)
        let candidate = PromptCandidate(
            hostProcess: hostProcess,
            texts: snapshot.texts,
            buttonTitles: snapshot.buttons.map(\.title)
        )

        if
            let match = AirPlayPromptMatcher.match(candidate),
            let acceptButton = snapshot.buttons.first(where: { $0.title == match.acceptButtonTitle })
        {
            return PromptCandidateWithButton(
                candidate: candidate,
                match: match,
                acceptButton: acceptButton.element,
                acceptActionName: acceptButton.actionName
            )
        }

        for child in SnapshotBuilder.children(of: element) {
            if let result = visit(
                child,
                hostProcess: hostProcess,
                depth: depth + 1,
                maxDepth: maxDepth,
                maxNodes: maxNodes,
                visited: &visited
            ) {
                return result
            }
        }

        return nil
    }
}

private struct SnapshotBuilder {
    struct Button {
        var title: String
        var element: AXUIElement
        var actionName: String
    }

    var texts: [String] = []
    var buttons: [Button] = []

    static func build(from root: AXUIElement, maxDepth: Int, maxNodes: Int) -> SnapshotBuilder {
        var builder = SnapshotBuilder()
        var visited = 0
        builder.visit(root, depth: 0, maxDepth: maxDepth, maxNodes: maxNodes, visited: &visited)
        return builder
    }

    private mutating func visit(
        _ element: AXUIElement,
        depth: Int,
        maxDepth: Int,
        maxNodes: Int,
        visited: inout Int
    ) {
        guard depth <= maxDepth, visited < maxNodes else {
            return
        }
        visited += 1

        let role = stringAttribute(element, kAXRoleAttribute)
        let title = stringAttribute(element, kAXTitleAttribute)
        let value = stringAttribute(element, kAXValueAttribute)
        let description = stringAttribute(element, kAXDescriptionAttribute)
        let help = stringAttribute(element, kAXHelpAttribute)
        let actionNames = actionNames(element)

        for text in [title, value, description, help].compactMap({ $0 }) {
            appendUnique(text, to: &texts)
        }

        if let button = buttonMetadata(
            role: role,
            title: title,
            value: value,
            description: description,
            help: help,
            actionNames: actionNames
        ) {
            buttons.append(Button(title: button.title, element: element, actionName: button.actionName))
        }

        for child in Self.children(of: element) {
            visit(child, depth: depth + 1, maxDepth: maxDepth, maxNodes: maxNodes, visited: &visited)
        }
    }

    private func stringAttribute(_ element: AXUIElement, _ attribute: String) -> String? {
        var value: CFTypeRef?
        guard AXUIElementCopyAttributeValue(element, attribute as CFString, &value) == .success else {
            return nil
        }

        if let string = value as? String {
            return string.trimmingCharacters(in: .whitespacesAndNewlines)
        }

        if let attributed = value as? NSAttributedString {
            return attributed.string.trimmingCharacters(in: .whitespacesAndNewlines)
        }

        return nil
    }

    static func children(of element: AXUIElement) -> [AXUIElement] {
        var value: CFTypeRef?
        guard
            AXUIElementCopyAttributeValue(element, kAXChildrenAttribute as CFString, &value) == .success,
            let children = value as? [AXUIElement]
        else {
            return []
        }
        return children
    }

    private func actionNames(_ element: AXUIElement) -> [String] {
        var value: CFArray?
        guard AXUIElementCopyActionNames(element, &value) == .success else {
            return []
        }
        return (value as? [CFString] ?? []).map { String($0) }
    }

    private func buttonMetadata(
        role: String?,
        title: String?,
        value: String?,
        description: String?,
        help: String?,
        actionNames: [String]
    ) -> (title: String, actionName: String)? {
        let labels = [title, value, description, help]
            .compactMap { $0?.trimmingCharacters(in: .whitespacesAndNewlines) }
            .filter { !$0.isEmpty }

        if let namedAction = actionNames.first(where: { actionTitle($0).map(AirPlayPromptMatcher.isAcceptTitle) == true }) {
            return (actionTitle(namedAction) ?? namedAction, namedAction)
        }

        if role == kAXButtonRole as String, let label = labels.first {
            return (label, kAXPressAction as String)
        }

        guard
            actionNames.contains(kAXPressAction as String),
            let label = labels.first(where: AirPlayPromptMatcher.isAcceptTitle)
        else {
            return nil
        }

        return (label, kAXPressAction as String)
    }

    private func actionTitle(_ actionName: String) -> String? {
        let normalized = actionName.trimmingCharacters(in: .whitespacesAndNewlines)
        guard normalized.localizedCaseInsensitiveContains("name:") else {
            return nil
        }

        let parts = normalized.split(separator: ":", maxSplits: 1)
        guard parts.count == 2 else {
            return nil
        }

        return String(parts[1]).trimmingCharacters(in: .whitespacesAndNewlines)
    }

    private func appendUnique(_ text: String, to array: inout [String]) {
        let cleaned = text.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !cleaned.isEmpty, !array.contains(cleaned) else {
            return
        }
        array.append(cleaned)
    }
}