import AppKit
import ApplicationServices

@_silgen_name("_AXUIElementGetWindow")
func _AXUIElementGetWindow(_ element: AXUIElement, _ id: UnsafeMutablePointer<CGWindowID>) -> AXError

struct RefSource {
    let pid: pid_t
    let window: Int?
}

var refs: [String: AXUIElement] = [:]
var refSource = RefSource(pid: 0, window: nil)

let shownActions: [String: String] = [
    kAXPressAction: "press",
    kAXConfirmAction: "confirm",
    kAXIncrementAction: "increment",
    kAXDecrementAction: "decrement",
    kAXPickAction: "pick",
    kAXCancelAction: "cancel",
]

let maxVisits = 30000

func attr(_ e: AXUIElement, _ name: String) -> CFTypeRef? {
    var v: CFTypeRef?
    if AXUIElementCopyAttributeValue(e, name as CFString, &v) != .success {
        return nil
    }
    return v
}

func text(_ e: AXUIElement, _ name: String) -> String {
    guard let v = attr(e, name) else {
        return ""
    }
    if let s = v as? String {
        return s
    }
    if CFGetTypeID(v) == CFBooleanGetTypeID() {
        return ""
    }
    if let n = v as? NSNumber {
        return n.stringValue
    }
    return ""
}

func children(_ e: AXUIElement) -> [AXUIElement] {
    return (attr(e, kAXChildrenAttribute) as? [AXUIElement]) ?? []
}

func appElement(_ pid: pid_t) -> AXUIElement {
    let app = AXUIElementCreateApplication(pid)
    AXUIElementSetMessagingTimeout(app, 3)
    return app
}

func axWindow(_ id: Int) throws -> (WindowInfo, AXUIElement) {
    let info = try window(id)
    for w in (attr(appElement(info.pid), kAXWindowsAttribute) as? [AXUIElement]) ?? [] {
        var wid: CGWindowID = 0
        if _AXUIElementGetWindow(w, &wid) == .success && Int(wid) == id {
            return (info, w)
        }
    }
    throw ToolError("window \(id) has no accessibility element; the app may not expose it")
}

func element(_ ref: String, in info: WindowInfo) throws -> AXUIElement {
    guard let e = refs[ref] else {
        throw ToolError("unknown ref \(ref); refs come from the latest ax_tree (or a tree returned by act) and are replaced by each new tree")
    }
    let s = refSource
    if s.pid != info.pid || (s.window != nil && s.window != Int(info.id)) {
        let from = s.window.map { "window \($0)" } ?? "the menu bar of process \(s.pid)"
        throw ToolError("\(ref) is from \(from), not window \(info.id); read window \(info.id)'s tree (or its app's menu bar) first")
    }
    return e
}

func frame(_ e: AXUIElement) -> CGRect? {
    guard let p = attr(e, kAXPositionAttribute), let s = attr(e, kAXSizeAttribute) else {
        return nil
    }
    var pt = CGPoint.zero
    var sz = CGSize.zero
    AXValueGetValue(p as! AXValue, .cgPoint, &pt)
    AXValueGetValue(s as! AXValue, .cgSize, &sz)
    return CGRect(origin: pt, size: sz)
}

func clip(_ s: String, _ n: Int) -> String {
    let flat = s.replacingOccurrences(of: "\n", with: " ")
    if flat.count <= n {
        return flat
    }
    return String(flat.prefix(n)) + "…"
}

struct TreeOptions {
    var query: String?
    var depth = 80
    var maxLines = 400
    var frames = false
    var origin = CGPoint.zero
}

func axTree(_ root: AXUIElement, _ source: RefSource, _ opts: TreeOptions) -> String {
    refs = [:]
    refSource = source
    var lines: [String] = []
    var visits = 0
    var truncated = false
    let q = opts.query?.lowercased()

    func walk(_ e: AXUIElement, _ depth: Int, _ indent: Int, _ parentParts: [String]) {
        visits += 1
        if visits > maxVisits || lines.count >= opts.maxLines {
            truncated = true
            return
        }
        let role = text(e, kAXRoleAttribute)
        var parts: [String] = []
        for name in [kAXTitleAttribute, kAXDescriptionAttribute, kAXValueAttribute, kAXPlaceholderValueAttribute] {
            let s = text(e, name)
            if !s.isEmpty && !parts.contains(s) {
                parts.append(s)
            }
        }
        var names: CFArray?
        AXUIElementCopyActionNames(e, &names)
        var acts = ((names as? [String]) ?? []).compactMap { shownActions[$0] }
        var settable: DarwinBoolean = false
        AXUIElementIsAttributeSettable(e, kAXValueAttribute as CFString, &settable)
        if settable.boolValue {
            acts.append("set")
        }
        if (attr(e, kAXFocusedAttribute) as? Bool) == true {
            acts.append("focused")
        }
        if (attr(e, kAXEnabledAttribute) as? Bool) == false {
            acts.append("disabled")
        }
        let echo = role == kAXStaticTextRole && acts.isEmpty && parts.allSatisfy { p in parentParts.contains { $0.contains(p) } }
        let shown = (!parts.isEmpty || !acts.isEmpty) && !echo
        let matches = q == nil || parts.contains { $0.lowercased().contains(q!) }
        var next = indent
        if shown && matches {
            let ref = "e\(refs.count + 1)"
            refs[ref] = e
            var line = String(repeating: "  ", count: q == nil ? indent : 0) + ref + " " + role.replacingOccurrences(of: "AX", with: "").lowercased()
            for p in parts {
                line += " " + clip(p, 200).debugDescription
            }
            if !acts.isEmpty {
                line += " [" + acts.joined(separator: " ") + "]"
            }
            if opts.frames, let f = frame(e) {
                line += " @\(Int(f.minX - opts.origin.x)),\(Int(f.minY - opts.origin.y)) \(Int(f.width))x\(Int(f.height))"
            }
            lines.append(line)
            next = indent + 1
        }
        if depth >= opts.depth {
            return
        }
        for c in children(e) {
            walk(c, depth + 1, next, parts.isEmpty ? parentParts : parts)
        }
    }

    walk(root, 0, 0, [])
    if truncated {
        lines.append("… truncated at \(lines.count) lines; narrow it with query or depth, or raise max_lines")
    }
    return lines.joined(separator: "\n")
}

func treeOptions(_ args: [String: Any], origin: CGPoint) -> TreeOptions {
    var o = TreeOptions()
    o.query = args.string("query")
    o.depth = args.int("depth") ?? o.depth
    o.maxLines = args.int("max_lines") ?? o.maxLines
    o.frames = args.bool("frames")
    o.origin = origin
    return o
}

func axTreeTool(_ args: [String: Any]) async throws -> [Content] {
    if let pid = args.int("menu_bar_pid") {
        guard let p = pid_t(exactly: pid), p > 0, let bar = attr(appElement(p), kAXMenuBarAttribute) else {
            throw ToolError("process \(pid) has no menu bar")
        }
        return [.text(axTree(bar as! AXUIElement, RefSource(pid: p, window: nil), treeOptions(args, origin: .zero)))]
    }
    let id = try args.require("window", args.int)
    let (info, w) = try axWindow(id)
    return [.text(axTree(w, RefSource(pid: info.pid, window: id), treeOptions(args, origin: info.frame.origin)))]
}

func windowText(_ root: AXUIElement, contains needle: String) -> Bool {
    let n = needle.lowercased()
    var visits = 0
    func walk(_ e: AXUIElement) -> Bool {
        visits += 1
        if visits > maxVisits {
            return false
        }
        for name in [kAXTitleAttribute, kAXDescriptionAttribute, kAXValueAttribute] {
            if text(e, name).lowercased().contains(n) {
                return true
            }
        }
        return children(e).contains { walk($0) }
    }
    return walk(root)
}