import Cocoa
import ApplicationServices

func extractTextAndCursor(from elem: AXUIElement) -> (String, Int, Int)? {
    var roleRef: CFTypeRef?
    AXUIElementCopyAttributeValue(elem, kAXRoleAttribute as CFString, &roleRef)
    let role = (roleRef as? String) ?? ""
    
    var currentText = ""
    var valueRef: CFTypeRef?
    if AXUIElementCopyAttributeValue(elem, kAXValueAttribute as CFString, &valueRef) == .success,
       let val = valueRef {
        if let str = val as? String {
            currentText = str
        } else if let attrStr = val as? NSAttributedString {
            currentText = attrStr.string
        }
    }
    
    if currentText.isEmpty {
        var countRef: CFTypeRef?
        if AXUIElementCopyAttributeValue(elem, kAXNumberOfCharactersAttribute as CFString, &countRef) == .success,
           let count = countRef as? Int, count > 0 {
            var range = CFRange(location: 0, length: count)
            if let axRange = AXValueCreate(.cfRange, &range) {
                var stringRef: CFTypeRef?
                if AXUIElementCopyParameterizedAttributeValue(elem, kAXStringForRangeParameterizedAttribute as CFString, axRange, &stringRef) == .success,
                   let str = stringRef as? String {
                    currentText = str
                }
            }
        }
    }
    
    var cursor = currentText.count
    var selLen = 0
    var rangeRef: CFTypeRef?
    if AXUIElementCopyAttributeValue(elem, kAXSelectedTextRangeAttribute as CFString, &rangeRef) == .success,
       let val = rangeRef {
        var cfRange = CFRange()
        if AXValueGetValue(val as! AXValue, .cfRange, &cfRange) {
            cursor = cfRange.location
            selLen = cfRange.length
        }
    }
    
    if !currentText.isEmpty || role == "AXTextField" || role == "AXTextArea" || role == "AXSearchField" || rangeRef != nil {
        return (currentText, cursor, selLen)
    }
    
    // Check focused child
    var focusedChildRef: CFTypeRef?
    if AXUIElementCopyAttributeValue(elem, kAXFocusedUIElementAttribute as CFString, &focusedChildRef) == .success,
       let child = focusedChildRef {
        if let res = extractTextAndCursor(from: child as! AXUIElement) {
            return res
        }
    }
    
    // Check children
    var childrenRef: CFTypeRef?
    if AXUIElementCopyAttributeValue(elem, kAXChildrenAttribute as CFString, &childrenRef) == .success,
       let children = childrenRef as? [AXUIElement] {
        for child in children {
            var isFocusedRef: CFTypeRef?
            if AXUIElementCopyAttributeValue(child, kAXFocusedAttribute as CFString, &isFocusedRef) == .success,
               let isFoc = isFocusedRef as? Bool, isFoc {
                if let res = extractTextAndCursor(from: child) {
                    return res
                }
            }
        }
    }
    
    return nil
}

func testFullAX() {
    let sysWide = AXUIElementCreateSystemWide()
    AXUIElementSetAttributeValue(sysWide, "AXEnhancedUserInterface" as CFString, kCFBooleanTrue)
    AXUIElementSetAttributeValue(sysWide, "AXManualAccessibility" as CFString, kCFBooleanTrue)
    
    guard let frontApp = NSWorkspace.shared.frontmostApplication else { return }
    print("Front App:", frontApp.localizedName ?? "")
    
    let appElem = AXUIElementCreateApplication(frontApp.processIdentifier)
    AXUIElementSetAttributeValue(appElem, "AXEnhancedUserInterface" as CFString, kCFBooleanTrue)
    AXUIElementSetAttributeValue(appElem, "AXManualAccessibility" as CFString, kCFBooleanTrue)
    
    var focusedElem: CFTypeRef?
    if AXUIElementCopyAttributeValue(sysWide, kAXFocusedUIElementAttribute as CFString, &focusedElem) == .success,
       let elem = focusedElem {
        if let (text, cursor, sel) = extractTextAndCursor(from: elem as! AXUIElement) {
            print("Extracted from sysWide -> Text: '\(text)', Cursor: \(cursor), Sel: \(sel)")
            return
        }
    }
    
    if AXUIElementCopyAttributeValue(appElem, kAXFocusedUIElementAttribute as CFString, &focusedElem) == .success,
       let elem = focusedElem {
        if let (text, cursor, sel) = extractTextAndCursor(from: elem as! AXUIElement) {
            print("Extracted from appElem -> Text: '\(text)', Cursor: \(cursor), Sel: \(sel)")
            return
        }
    }
    
    print("Could not extract text")
}

testFullAX()