import AppKit
import ApplicationServices
import Foundation

/// Per-connection registry of the element handles (`ref` → live `AXUIElementRef`) issued by the
/// last `ax_dump`/`ax_element_at` (§6.2). `AXUIElementRef`s go stale as the UI mutates, so refs are
/// valid until the next dump replaces them; a stale ref surfaces as "re-dump and retry".
final class AXRegistry {
    private var refs: [String: AXUIElement] = [:]
    private var counter = 0

    func clear() { refs.removeAll() }

    func register(_ element: AXUIElement) -> String {
        counter += 1
        let ref = "e\(counter)"
        refs[ref] = element
        return ref
    }

    subscript(ref: String) -> AXUIElement? { refs[ref] }
}

/// Accessibility observe + control (§6): the semantic, framebuffer-independent surface. All calls
/// are synchronous cross-process IPC into the target app, bounded by the messaging timeout so a
/// hung app can't wedge the walker.
extension ConnectionHandler {
    private var axTrusted: Bool { AXIsProcessTrusted() }
    private static let axDenied =
        "Accessibility (kTCCServiceAccessibility) is not granted to the agent — re-run provisioning"

    /// `ax_dump` — walk the frontmost (or `pid`-addressed) app's element tree into a JSON node tree,
    /// registering a fresh `ref` per element. This is the model's semantic "look" primitive.
    func opAXDump(_ request: [String: Any]) -> [String: Any] {
        guard axTrusted else { return AgentReply.err(Self.axDenied) }
        let pid: pid_t
        if let requested = request["pid"] as? Int {
            pid = pid_t(requested)
        } else if let front = NSWorkspace.shared.frontmostApplication {
            pid = front.processIdentifier
        } else {
            return AgentReply.err("no frontmost application to dump")
        }
        let maxDepth = min(max((request["maxDepth"] as? Int) ?? 12, 1), 25)

        // New dump ⇒ new handle generation; every previously issued ref is now invalid by contract.
        axRegistry.clear()
        let app = AXUIElementCreateApplication(pid)
        // A tight timeout so ONE hung app can't stall the whole walk for the default ~6 s per call.
        AXUIElementSetMessagingTimeout(app, 1.5)

        // Total node budget: a rich app's full tree is thousands of IPC round-trips (§6.1); the cap
        // bounds both the walk time and the payload the model has to read.
        var budget = 800
        let tree = buildNode(app, depth: 0, maxDepth: maxDepth, budget: &budget)
        var fields: [String: Any] = ["tree": tree]
        fields["app"] = NSRunningApplication(processIdentifier: pid)?.localizedName ?? "pid \(pid)"
        if budget <= 0 { fields["truncated"] = true }
        return AgentReply.ok(fields)
    }

    /// `ax_element_at` — resolve the element under a screen point (hit-test via the system-wide
    /// element) into a single node (no children). Registers a ref alongside the last dump's.
    func opAXElementAt(_ request: [String: Any]) -> [String: Any] {
        guard axTrusted else { return AgentReply.err(Self.axDenied) }
        guard let x = request["x"] as? Int, let y = request["y"] as? Int else {
            return AgentReply.err("ax_element_at requires x and y")
        }
        let systemWide = AXUIElementCreateSystemWide()
        AXUIElementSetMessagingTimeout(systemWide, 1.5)
        var element: AXUIElement?
        let result = AXUIElementCopyElementAtPosition(systemWide, Float(x), Float(y), &element)
        guard result == .success, let element else {
            return AgentReply.err("no accessible element at (\(x), \(y)) (\(describe(result)))")
        }
        var budget = 1
        return AgentReply.ok(["node": buildNode(element, depth: 0, maxDepth: 0, budget: &budget)])
    }

    /// `ax_action` — perform a named action (`AXPress`, `AXShowMenu`, …) on a ref'd element: the
    /// control's own callback runs in the app process, no rendering or synthetic click involved.
    func opAXAction(_ request: [String: Any]) -> [String: Any] {
        guard axTrusted else { return AgentReply.err(Self.axDenied) }
        guard let element = resolveRef(request) else { return staleRefReply(request) }
        let action = (request["action"] as? String).flatMap { $0.isEmpty ? nil : $0 } ?? "AXPress"
        let result = AXUIElementPerformAction(element, action as CFString)
        switch result {
        case .success:
            return AgentReply.ok()
        case .cannotComplete:
            // Apps often run modal processing inside the action callback and never return within
            // the AX timeout — the action may well have fired (§6.3). Tell the host the truth.
            return AgentReply.ok([
                "note": "the app did not confirm within the AX timeout — the action possibly "
                    + "succeeded; re-dump to verify"
            ])
        case .invalidUIElement:
            return staleRefReply(request)
        default:
            return AgentReply.err("\(action) failed (\(describe(result)))")
        }
    }

    /// `ax_set_value` — write a field's `AXValue` directly (the reliable way to fill text fields).
    /// Non-settable / non-AppKit fields are reported as such so the host can fall back to
    /// `ax_focus` + CGEvent typing.
    func opAXSetValue(_ request: [String: Any]) -> [String: Any] {
        guard axTrusted else { return AgentReply.err(Self.axDenied) }
        guard let element = resolveRef(request) else { return staleRefReply(request) }
        guard let value = request["value"] as? String else {
            return AgentReply.err("ax_set_value requires a string \"value\"")
        }
        var settable = DarwinBoolean(false)
        AXUIElementIsAttributeSettable(element, kAXValueAttribute as CFString, &settable)
        guard settable.boolValue else {
            return AgentReply.err(
                "this element's AXValue is not settable — focus it (ax_focus) and type instead")
        }
        let result = AXUIElementSetAttributeValue(
            element, kAXValueAttribute as CFString, value as CFString)
        guard result == .success else {
            if result == .invalidUIElement { return staleRefReply(request) }
            return AgentReply.err("setting AXValue failed (\(describe(result)))")
        }
        return AgentReply.ok()
    }

    /// `ax_focus` — give a ref'd element keyboard focus (precedes CGEvent typing into it).
    func opAXFocus(_ request: [String: Any]) -> [String: Any] {
        guard axTrusted else { return AgentReply.err(Self.axDenied) }
        guard let element = resolveRef(request) else { return staleRefReply(request) }
        let result = AXUIElementSetAttributeValue(
            element, kAXFocusedAttribute as CFString, kCFBooleanTrue)
        guard result == .success else {
            if result == .invalidUIElement { return staleRefReply(request) }
            return AgentReply.err("focusing failed (\(describe(result)))")
        }
        return AgentReply.ok()
    }

    // MARK: - Tree walking

    /// Serialize one element into the §4 node shape, recursing into children while depth and the
    /// total node budget allow. Empty/absent attributes are omitted to keep the payload the model
    /// reads small.
    private func buildNode(
        _ element: AXUIElement, depth: Int, maxDepth: Int, budget: inout Int
    ) -> [String: Any] {
        budget -= 1
        var node: [String: Any] = ["ref": axRegistry.register(element)]
        if let role = copyString(element, kAXRoleAttribute) { node["role"] = role }
        if let subrole = copyString(element, kAXSubroleAttribute) { node["subrole"] = subrole }
        if let title = copyString(element, kAXTitleAttribute), !title.isEmpty {
            node["title"] = title
        }
        if let value = copyValueDescription(element) { node["value"] = value }
        if copyBool(element, kAXEnabledAttribute) == false { node["enabled"] = false }
        if copyBool(element, kAXFocusedAttribute) == true { node["focused"] = true }
        if let frame = copyFrame(element) { node["frame"] = frame }

        var actionNames: CFArray?
        if AXUIElementCopyActionNames(element, &actionNames) == .success,
            let actions = actionNames as? [String], !actions.isEmpty
        {
            node["actions"] = actions
        }

        if depth < maxDepth, budget > 0 {
            var childrenRef: CFTypeRef?
            if AXUIElementCopyAttributeValue(
                element, kAXChildrenAttribute as CFString, &childrenRef) == .success,
                let children = childrenRef as? [AXUIElement], !children.isEmpty
            {
                var serialized: [[String: Any]] = []
                // Per-node child cap: a 10k-row table must not eat the entire budget at one level.
                for child in children.prefix(48) {
                    guard budget > 0 else { break }
                    serialized.append(
                        buildNode(child, depth: depth + 1, maxDepth: maxDepth, budget: &budget))
                }
                if !serialized.isEmpty { node["children"] = serialized }
                if children.count > serialized.count {
                    node["childrenOmitted"] = children.count - serialized.count
                }
            }
        }
        return node
    }

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

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

    /// A human/model-readable rendering of `AXValue` (string, number, bool), truncated so one giant
    /// text view can't dominate the tree payload.
    private func copyValueDescription(_ element: AXUIElement) -> String? {
        var value: CFTypeRef?
        guard
            AXUIElementCopyAttributeValue(element, kAXValueAttribute as CFString, &value) == .success,
            let value
        else { return nil }
        let text: String
        if let string = value as? String {
            text = string
        } else if let number = value as? NSNumber {
            text = number.stringValue
        } else if let url = value as? URL {
            text = url.absoluteString
        } else {
            return nil  // AXValueRef boxes (ranges etc.) aren't useful as text
        }
        if text.isEmpty { return nil }
        return text.count > 250 ? String(text.prefix(250)) + "…" : text
    }

    /// `kAXPosition`/`kAXSize` come back as `AXValueRef` boxes — unwrap into the §4 frame shape
    /// (top-left-origin screen coordinates, matching screenshots and CGEvent clicks).
    private func copyFrame(_ element: AXUIElement) -> [String: Any]? {
        var positionRef: CFTypeRef?
        var sizeRef: CFTypeRef?
        var point = CGPoint.zero
        var size = CGSize.zero
        guard
            AXUIElementCopyAttributeValue(
                element, kAXPositionAttribute as CFString, &positionRef) == .success,
            AXUIElementCopyAttributeValue(
                element, kAXSizeAttribute as CFString, &sizeRef) == .success,
            let positionRef, let sizeRef,
            AXValueGetValue(positionRef as! AXValue, .cgPoint, &point),
            AXValueGetValue(sizeRef as! AXValue, .cgSize, &size)
        else { return nil }
        return [
            "x": Int(point.x.rounded()), "y": Int(point.y.rounded()),
            "w": Int(size.width.rounded()), "h": Int(size.height.rounded()),
        ]
    }

    // MARK: - Ref plumbing

    private func resolveRef(_ request: [String: Any]) -> AXUIElement? {
        guard let ref = request["ref"] as? String else { return nil }
        return axRegistry[ref]
    }

    private func staleRefReply(_ request: [String: Any]) -> [String: Any] {
        let ref = (request["ref"] as? String) ?? "?"
        return AgentReply.err(
            "element ref \"\(ref)\" is unknown or stale — refs are only valid until the next "
                + "ax_dump; re-dump and retry")
    }

    private func describe(_ error: AXError) -> String {
        switch error {
        case .success: return "success"
        case .apiDisabled: return "kAXErrorAPIDisabled — accessibility is disabled for this app"
        case .cannotComplete: return "kAXErrorCannotComplete — the app did not respond in time"
        case .invalidUIElement: return "kAXErrorInvalidUIElement — stale element"
        case .attributeUnsupported: return "attribute unsupported"
        case .actionUnsupported: return "action unsupported"
        case .noValue: return "no value"
        case .notImplemented: return "the app does not implement this AX request"
        default: return "AXError \(error.rawValue)"
        }
    }
}
