import AppKit
import ApplicationServices

func point(_ v: Any?, _ info: WindowInfo, _ table: RefTable) throws -> CGPoint {
    if let ref = v as? String {
        guard let f = frame(try element(ref, in: info, table)) else {
            throw ToolError("\(ref) has no position on screen")
        }
        return CGPoint(x: f.midX, y: f.midY)
    }
    if let xy = v as? [NSNumber], xy.count == 2 {
        return CGPoint(x: info.frame.minX + CGFloat(xy[0].doubleValue), y: info.frame.minY + CGFloat(xy[1].doubleValue))
    }
    throw ToolError("a point is a ref (\"e5\") or window coordinates ([x, y])")
}

func check(_ err: AXError, _ what: String) throws {
    if err != .success {
        throw ToolError("\(what) failed: AXError \(err.rawValue)")
    }
}

let inputKeys = ["click", "double_click", "right_click", "drag", "scroll", "type", "key"]

func isInput(_ step: [String: Any]) -> Bool {
    return step.keys.contains(where: inputKeys.contains)
}

func runStep(_ step: [String: Any], _ id: Int, _ w: AXUIElement, _ table: RefTable) async throws -> String {
    var info = try window(id)
    if isInput(step) {
        info = try await focusWindow(id, w)
    }
    if let ref = step.string("press") {
        try check(AXUIElementPerformAction(try element(ref, in: info, table), kAXPressAction as CFString), "press \(ref)")
        return "pressed \(ref)"
    }
    if let ref = step.string("confirm") {
        try check(AXUIElementPerformAction(try element(ref, in: info, table), kAXConfirmAction as CFString), "confirm \(ref)")
        return "confirmed \(ref)"
    }
    if let ref = step.string("focus") {
        try check(AXUIElementSetAttributeValue(try element(ref, in: info, table), kAXFocusedAttribute as CFString, kCFBooleanTrue), "focus \(ref)")
        return "focused \(ref)"
    }
    if let ref = step.string("set") {
        let value = try step.require("value", step.string)
        try check(AXUIElementSetAttributeValue(try element(ref, in: info, table), kAXValueAttribute as CFString, value as CFString), "set \(ref)")
        return "set \(ref)"
    }
    if let name = step.string("action") {
        let ref = try step.require("ref", step.string)
        try check(AXUIElementPerformAction(try element(ref, in: info, table), name as CFString), "\(name) \(ref)")
        return "\(name) \(ref)"
    }
    for (k, right, count) in [("click", false, 1), ("double_click", false, 2), ("right_click", true, 1)] {
        if step[k] != nil {
            let p = try point(step[k], info, table)
            await click(p, right: right, count: count)
            return "\(k) at screen \(Int(p.x)),\(Int(p.y))"
        }
    }
    if step["drag"] != nil {
        let from = try point(step["drag"], info, table)
        let to = try point(step["to"], info, table)
        await drag(from, to)
        return "dragged"
    }
    if step["scroll"] != nil {
        let p = try point(step["scroll"], info, table)
        let dx = step.int("dx") ?? 0
        let dy = step.int("dy") ?? 0
        guard (-1000...1000).contains(dx), (-1000...1000).contains(dy) else {
            throw ToolError("dx and dy must be between -1000 and 1000 lines")
        }
        scroll(p, dx: Int32(dx), dy: Int32(dy))
        return "scrolled"
    }
    if let s = step.string("type") {
        await type(s)
        return "typed \(s.count) characters"
    }
    if let k = step.string("key") {
        try key(k)
        return "key \(k)"
    }
    if let ms = try step.ms("wait") {
        await pause(ms)
        return "waited \(ms)ms"
    }
    throw ToolError("unknown step \(step.keys.sorted())")
}

func waitFor(_ args: [String: Any], acted: Bool, _ present: (String) async throws -> Bool) async throws -> String? {
    let wait = try args.ms("wait_ms")
    let timeout = Double(try args.ms("timeout_ms") ?? 10000) / 1000
    if let wait {
        await pause(wait)
    } else if acted {
        await pause(300)
    }
    guard let needle = args.string("wait_for") else {
        return nil
    }
    let start = Date()
    while true {
        if try await present(needle) {
            return String(format: "found %@ after %.1fs", needle.debugDescription, Date().timeIntervalSince(start))
        }
        if Date().timeIntervalSince(start) > timeout {
            return String(format: "gave up waiting for %@ after %.1fs", needle.debugDescription, timeout)
        }
        await pause(200)
    }
}

func runSteps(_ steps: [[String: Any]], _ id: Int, _ w: AXUIElement, _ table: RefTable) async throws -> [String] {
    var log: [String] = []
    for (i, step) in steps.enumerated() {
        do {
            log.append("\(i + 1). " + (try await runStep(step, id, w, table)))
            await pause(50)
        } catch let e as ToolError {
            log.append("\(i + 1). " + e.message + "; stopped here")
            break
        }
    }
    return log
}

func actTool(_ args: [String: Any]) async throws -> [Content] {
    let table = currentRefs()
    let id = try args.require("window", args.int)
    let (_, w) = try axWindow(id)
    let steps = args["steps"] as? [[String: Any]] ?? []
    let lock = steps.isEmpty ? nil : try await lockInput()
    defer {
        if let lock {
            close(lock)
        }
    }
    var log = try await runSteps(steps, id, w, table)
    if let note = try await waitFor(args, acted: !steps.isEmpty, { try windowText(w, contains: $0) }) {
        log.append(note)
    }
    var out: [Content] = []
    if args.bool("tree") {
        let current = try window(id)
        log.append("\n" + (try axTree(w, RefSource(pid: current.pid, window: id), treeOptions(args, origin: current.frame.origin))))
    }
    if args.bool("capture") {
        let shot = try await capture(window: id, display: nil, crop: args.dict("crop"), maxEdge: args.double("max_edge") ?? 1568)
        log.append(shot.note)
        out.append(.png(shot.png))
    }
    return [.text(log.joined(separator: "\n"))] + out
}