import Foundation
import NucleicProtocol

/// Computer-use over the **native in-guest agent** (docs/MACOS_VM_NATIVE_AGENT.md): the preferred
/// path when the guest runs `NucleicVMAgent` — ScreenCaptureKit capture, CGEvent input (with a REAL
/// scroll wheel), and the AX-semantic ops that have no SSH/cliclick equivalent at all. The SSH path
/// (`MacVMEngine+Computer`) remains the transparent fallback, so a base image without the agent
/// loses nothing.
extension MacVMEngine {
    /// The AX-semantic tool actions (§10). Agent-only: the accessibility tree is not reachable over
    /// the SSH/cliclick path, so these never fall back.
    static let axActions: Set<String> = [
        "ax_dump", "ax_element_at", "ax_press", "ax_set_value", "ax_focus",
    ]

    /// Try one pixel/input action natively. Returns the reply when the action ran (never falls
    /// back past this point — re-running a click over SSH would double-click), or `nil` when the
    /// action did NOT run (op failed before side effects, or the args don't validate) and the SSH
    /// path should take over.
    func performAgentComputerAction(
        name: String, client: MacVMAgentClient, action: String, x: Int?, y: Int?, text: String?,
        scrollDirection: String?, scrollAmount: Int?, durationMs: Int?
    ) async -> (imageBase64: String?, summary: String)? {
        do {
            switch action {
            case "screenshot":
                break  // nothing to do but capture below

            case "cursor_position":
                let reply = try await client.request(op: "cursor_position")
                guard let cx = reply["x"]?.intValue, let cy = reply["y"]?.intValue else {
                    return nil
                }
                return (nil, "Cursor position: \(cx),\(cy)")

            case "wait":
                let ms = max(0, min(durationMs ?? 1000, 30_000))
                try? await Task.sleep(nanoseconds: UInt64(ms) * 1_000_000)

            case "left_click", "right_click", "double_click":
                guard let x, let y else { return nil }
                var fields: [String: JSONValue] = [
                    "x": .number(Double(x)), "y": .number(Double(y)),
                ]
                if action == "right_click" { fields["button"] = .string("right") }
                if action == "double_click" { fields["count"] = .number(2) }
                _ = try await client.request(op: "click", fields: fields)

            case "mouse_move":
                guard let x, let y else { return nil }
                _ = try await client.request(
                    op: "move", fields: ["x": .number(Double(x)), "y": .number(Double(y))])

            case "left_click_drag":
                // Drag from the CURRENT cursor to the target (the agent defaults `from` to the
                // live pointer, so no cursor-position round-trip is needed).
                guard let x, let y else { return nil }
                _ = try await client.request(
                    op: "drag", fields: ["toX": .number(Double(x)), "toY": .number(Double(y))])

            case "type":
                guard let text, !text.isEmpty else { return nil }
                _ = try await client.request(op: "type", fields: ["text": .string(text)])

            case "key":
                guard let chord = text, !chord.isEmpty else { return nil }
                _ = try await client.request(op: "key", fields: ["chord": .string(chord)])

            case "scroll":
                let amount = max(1, min(scrollAmount ?? 3, 30))
                let (dx, dy) = Self.scrollDelta(
                    direction: scrollDirection ?? "down", amount: amount)
                // The agent paces the total into wheel notches itself, and moves the guest pointer to
                // `x`,`y` first when given (a wheel event routes by pointer, not by coordinate).
                var fields: [String: JSONValue] = [
                    "dx": .number(Double(dx)), "dy": .number(Double(dy)),
                ]
                if let x, let y {
                    fields["x"] = .number(Double(x))
                    fields["y"] = .number(Double(y))
                }
                _ = try await client.request(op: "scroll", fields: fields)

            case "launch_app":
                guard let app = text, !app.isEmpty else { return nil }
                _ = try await client.request(op: "launch_app", fields: ["name": .string(app)])

            default:
                return nil  // not a native action — the SSH path owns the "unknown action" reply
            }
        } catch {
            await noteAgentFailure(name: name, client: client, error: error)
            return nil  // the op failed before any side effect — safe to re-run over SSH
        }

        // Let the guest repaint before grabbing (async repaint, else the frame is pre-action). The
        // native screenshot is a vsock round-trip — too costly to poll for a change — so this uses one
        // fixed settle rather than the surface path's adaptive wait.
        if Self.actionNeedsSettle(action) {
            try? await Task.sleep(nanoseconds: Self.settleFixedNanos)
        }
        // The action ran natively; capture the post-action frame. Capture-only fallbacks from here
        // (the action itself must never re-run).
        if let shot = await agentScreenshot(name: name, client: client) {
            return (shot, screenSummary(for: action, name: name))
        }
        if let shot = try? await captureScreenshotBase64(name: name) {
            return (shot, screenSummary(for: action, name: name))
        }
        return (
            nil,
            screenSummary(for: action, name: name) + " No screenshot came back — the guest may lack Screen "
                + "Recording permission, or the capture may be unavailable. Use `ax_dump` to inspect "
                + "the UI semantically instead."
        )
    }

    /// Direction/amount (the tool's scroll vocabulary) → wheel-line deltas (the agent's §4 `scroll`
    /// op; dy > 0 scrolls content up, dx > 0 scrolls left — CGEvent wheel conventions).
    static func scrollDelta(direction: String, amount: Int) -> (dx: Int, dy: Int) {
        switch direction.lowercased() {
        case "up": return (0, amount)
        case "left": return (amount, 0)
        case "right": return (-amount, 0)
        default: return (0, -amount)  // down
        }
    }

    // MARK: - AX-semantic actions (§6, §10)

    /// Perform one AX action. Unlike the pixel path this never falls back to SSH (there is no SSH
    /// equivalent), so failures come back as actionable text for the model.
    func performAXComputerAction(
        name: String, client: MacVMAgentClient, action: String, x: Int?, y: Int?, text: String?,
        ref: String?, value: String?
    ) async -> (imageBase64: String?, summary: String) {
        do {
            switch action {
            case "ax_dump":
                let reply = try await client.request(op: "ax_dump")
                let app = reply["app"]?.stringValue ?? "frontmost app"
                let tree = reply["tree"] ?? .object([:])
                var header = "Accessibility tree of \(app) (roles/titles/values/frames; act on an "
                    + "element by its `ref` with ax_press / ax_set_value / ax_focus — refs stay "
                    + "valid until the next ax_dump):"
                if reply["truncated"]?.boolValue == true {
                    header += " [tree truncated — deep/rich UI]"
                }
                return (nil, header + "\n" + tree.canonicalString())

            case "ax_element_at":
                guard let x, let y else {
                    return (nil, "ax_element_at requires x and y coordinates.")
                }
                let reply = try await client.request(
                    op: "ax_element_at",
                    fields: ["x": .number(Double(x)), "y": .number(Double(y))])
                let node = reply["node"] ?? .object([:])
                return (nil, "Element at (\(x), \(y)): " + node.canonicalString())

            case "ax_press":
                guard let ref, !ref.isEmpty else {
                    return (nil, "ax_press requires an element `ref` (from ax_dump).")
                }
                var fields: [String: JSONValue] = ["ref": .string(ref)]
                // `text` optionally names a non-default AX action (AXShowMenu, AXConfirm, …).
                if let axAction = text, !axAction.isEmpty { fields["action"] = .string(axAction) }
                let reply = try await client.request(op: "ax_action", fields: fields)
                // Coordinate fallback (Linux Phase 2): an element with no semantic action comes back
                // with `fallback:"click"` + a true global `frame` (window origin from the compositor
                // helper + AT-SPI window-relative extents). Click its center via the host surface —
                // the pixel path that's 1:1 with the screen — instead of dead-ending.
                if reply["fallback"]?.stringValue == "click", let frame = reply["frame"]?.objectValue,
                    let fx = frame["x"]?.intValue, let fy = frame["y"]?.intValue,
                    let fw = frame["w"]?.intValue, let fh = frame["h"]?.intValue, fw > 0, fh > 0,
                    hasComputerSurface(name)
                {
                    let cx = fx + fw / 2, cy = fy + fh / 2
                    return await performSurfaceComputerAction(
                        name: name, action: "left_click", x: cx, y: cy, text: nil,
                        scrollDirection: nil, scrollAmount: nil, durationMs: nil)
                }
                var summary = "Performed \(text?.isEmpty == false ? text! : "AXPress") on \(ref)."
                if let note = reply["note"]?.stringValue { summary += " Note: \(note)." }
                return await withPostActionScreenshot(
                    name: name, client: client, summary: summary)

            case "ax_set_value":
                guard let ref, !ref.isEmpty else {
                    return (nil, "ax_set_value requires an element `ref` (from ax_dump).")
                }
                guard let value else {
                    return (nil, "ax_set_value requires `value` (the text to set).")
                }
                _ = try await client.request(
                    op: "ax_set_value",
                    fields: ["ref": .string(ref), "value": .string(value)])
                return await withPostActionScreenshot(
                    name: name, client: client, summary: "Set the value of \(ref).")

            case "ax_focus":
                guard let ref, !ref.isEmpty else {
                    return (nil, "ax_focus requires an element `ref` (from ax_dump).")
                }
                _ = try await client.request(op: "ax_focus", fields: ["ref": .string(ref)])
                return await withPostActionScreenshot(
                    name: name, client: client, summary: "Focused \(ref) — you can now `type` into it.")

            default:
                return (nil, "Unknown AX action '\(action)'.")
            }
        } catch {
            await noteAgentFailure(name: name, client: client, error: error)
            let reason = (error as? MacVMAgentError)?.description ?? String(describing: error)
            return (
                nil,
                "\(action) failed: \(reason) If the element may have changed, re-run ax_dump and "
                    + "retry with a fresh ref; otherwise fall back to screenshot + pixel actions."
            )
        }
    }

    /// The mutating AX actions return a screenshot when capture works (helpful on healthy guests)
    /// and a re-dump hint always (the ONLY observation channel when capture is blank or unavailable).
    private func withPostActionScreenshot(
        name: String, client: MacVMAgentClient, summary: String
    ) async -> (imageBase64: String?, summary: String) {
        let hint = " Re-run ax_dump to observe the result semantically (refs may now be stale)."
        if let shot = await agentScreenshot(name: name, client: client) {
            return (shot, summary + hint)
        }
        return (nil, summary + hint)
    }

    /// One native screenshot, or `nil` (after noting a transport failure). Never throws — capture
    /// is always best-effort around an action that already ran.
    private func agentScreenshot(name: String, client: MacVMAgentClient) async -> String? {
        do {
            let reply = try await client.request(op: "screenshot")
            return reply["image"]?.stringValue
        } catch {
            await noteAgentFailure(name: name, client: client, error: error)
            return nil
        }
    }

    /// Bookkeeping for a failed agent request: a transport failure discards the channel (the next
    /// action re-probes once); an op-level `agentError` keeps the healthy channel. Internal (not
    /// private) because the long-text typing route in `MacVMEngine+Computer` shares it.
    func noteAgentFailure(name: String, client: MacVMAgentClient, error: Error) async {
        if let agentError = error as? MacVMAgentError, agentError.isTransportFailure {
            await discardAgentClient(name: name, client)
        }
    }
}
