"""macgic_mcp — macOS Accessibility MCP Server.

Exposes macOS Accessibility and Core Graphics APIs as MCP tools,
giving AI systems full control of a Mac: clicking, typing, drawing,
reading UI, launching apps, and taking screenshots.
"""

import base64
import time

import AppKit
import ApplicationServices as AS
import Quartz
from mcp.server.fastmcp import FastMCP, Image

mcp = FastMCP("macgic")

# ---------------------------------------------------------------------------
# Key code mappings
# ---------------------------------------------------------------------------

KEY_CODES: dict[str, int] = {
    "a": 0, "b": 11, "c": 8, "d": 2, "e": 14, "f": 3, "g": 5, "h": 4,
    "i": 34, "j": 38, "k": 40, "l": 37, "m": 46, "n": 45, "o": 31, "p": 35,
    "q": 12, "r": 15, "s": 1, "t": 17, "u": 32, "v": 9, "w": 13, "x": 7,
    "y": 16, "z": 6,
    "0": 29, "1": 18, "2": 19, "3": 20, "4": 21, "5": 23, "6": 22, "7": 26,
    "8": 28, "9": 25,
    "return": 36, "enter": 36, "tab": 48, "space": 49, "delete": 51,
    "backspace": 51, "escape": 53, "esc": 53,
    "up": 126, "down": 125, "left": 123, "right": 124,
    "f1": 122, "f2": 120, "f3": 99, "f4": 118, "f5": 96, "f6": 97,
    "f7": 98, "f8": 100, "f9": 101, "f10": 109, "f11": 103, "f12": 111,
    "-": 27, "=": 24, "[": 33, "]": 30, "\\": 42, ";": 41, "'": 39,
    ",": 43, ".": 47, "/": 44, "`": 50,
    "home": 115, "end": 119, "pageup": 116, "pagedown": 121,
    "forwarddelete": 117,
}

# Characters that require Shift to type
SHIFT_CHARS: dict[str, str] = {
    "!": "1", "@": "2", "#": "3", "$": "4", "%": "5", "^": "6", "&": "7",
    "*": "8", "(": "9", ")": "0", "_": "-", "+": "=", "{": "[", "}": "]",
    "|": "\\", ":": ";", '"': "'", "<": ",", ">": ".", "?": "/", "~": "`",
}

MODIFIER_FLAGS: dict[str, int] = {
    "cmd": Quartz.kCGEventFlagMaskCommand,
    "command": Quartz.kCGEventFlagMaskCommand,
    "shift": Quartz.kCGEventFlagMaskShift,
    "alt": Quartz.kCGEventFlagMaskAlternate,
    "option": Quartz.kCGEventFlagMaskAlternate,
    "opt": Quartz.kCGEventFlagMaskAlternate,
    "ctrl": Quartz.kCGEventFlagMaskControl,
    "control": Quartz.kCGEventFlagMaskControl,
    "fn": Quartz.kCGEventFlagMaskSecondaryFn,
}

# ---------------------------------------------------------------------------
# Vision / Context tools
# ---------------------------------------------------------------------------


@mcp.tool()
def screenshot(
    display_id: int | None = None,
    x: float | None = None,
    y: float | None = None,
    width: float | None = None,
    height: float | None = None,
) -> Image:
    """Capture a screenshot. Optionally specify a display or a region (x, y, width, height)."""
    if x is not None and y is not None and width is not None and height is not None:
        rect = Quartz.CGRectMake(x, y, width, height)
    elif display_id is not None:
        rect = Quartz.CGDisplayBounds(display_id)
    else:
        rect = Quartz.CGRectInfinite

    image = Quartz.CGWindowListCreateImage(
        rect,
        Quartz.kCGWindowListOptionOnScreenOnly,
        Quartz.kCGNullWindowID,
        Quartz.kCGWindowImageDefault,
    )
    if image is None:
        raise RuntimeError(
            "Failed to capture screenshot. Ensure Screen Recording permission is granted."
        )

    # Convert CGImage → PNG bytes
    bitmap_rep = AppKit.NSBitmapImageRep.alloc().initWithCGImage_(image)
    png_data = bitmap_rep.representationUsingType_properties_(
        AppKit.NSBitmapImageFileTypePNG, {}
    )
    return Image(data=base64.b64decode(base64.b64encode(bytes(png_data))), format="png")


@mcp.tool()
def list_displays() -> list[dict]:
    """List all connected displays with IDs, resolutions, and positions."""
    max_displays = 16
    (err, display_ids, count) = Quartz.CGGetActiveDisplayList(max_displays, None, None)
    if err != 0:
        raise RuntimeError(f"CGGetActiveDisplayList failed with error {err}")

    displays = []
    for did in display_ids[:count]:
        bounds = Quartz.CGDisplayBounds(did)
        displays.append({
            "display_id": did,
            "width": int(bounds.size.width),
            "height": int(bounds.size.height),
            "origin_x": int(bounds.origin.x),
            "origin_y": int(bounds.origin.y),
            "is_main": bool(Quartz.CGDisplayIsMain(did)),
        })
    return displays


@mcp.tool()
def list_apps(include_background: bool = False) -> list[dict]:
    """List running applications (name, bundle ID, PID).

    By default only shows regular (visible) apps. Set include_background=True
    to also include background-only apps.
    """
    workspace = AppKit.NSWorkspace.sharedWorkspace()
    apps = workspace.runningApplications()
    result = []
    for app in apps:
        if (
            not include_background
            and app.activationPolicy() != AppKit.NSApplicationActivationPolicyRegular
        ):
            continue
        result.append({
            "name": str(app.localizedName() or ""),
            "bundle_id": str(app.bundleIdentifier() or ""),
            "pid": app.processIdentifier(),
            "is_active": bool(app.isActive()),
        })
    return result


@mcp.tool()
def get_ui_elements(pid: int, max_depth: int = 3) -> list[dict]:
    """Read the accessibility tree for an app identified by PID.

    Returns elements with role, title, value, position, size, and enabled state.
    max_depth controls how deep to recurse (default 3).
    """
    app_ref = AS.AXUIElementCreateApplication(pid)

    def _walk(element, depth: int) -> list[dict]:
        if depth > max_depth:
            return []

        items: list[dict] = []
        entry: dict = {}

        # Role
        err, role = AS.AXUIElementCopyAttributeValue(element, "AXRole", None)
        if err == 0 and role:
            entry["role"] = str(role)

        # Title
        err, title = AS.AXUIElementCopyAttributeValue(element, "AXTitle", None)
        if err == 0 and title:
            entry["title"] = str(title)

        # Value
        err, value = AS.AXUIElementCopyAttributeValue(element, "AXValue", None)
        if err == 0 and value is not None:
            entry["value"] = str(value)

        # Description
        err, desc = AS.AXUIElementCopyAttributeValue(element, "AXDescription", None)
        if err == 0 and desc:
            entry["description"] = str(desc)

        # Enabled
        err, enabled = AS.AXUIElementCopyAttributeValue(element, "AXEnabled", None)
        if err == 0 and enabled is not None:
            entry["enabled"] = bool(enabled)

        # Position & Size
        err, pos_val = AS.AXUIElementCopyAttributeValue(element, "AXPosition", None)
        if err == 0 and pos_val is not None:
            ok, point = AS.AXValueGetValue(pos_val, AS.kAXValueTypeCGPoint, None)
            if ok:
                entry["x"] = round(point.x, 1)
                entry["y"] = round(point.y, 1)

        err, size_val = AS.AXUIElementCopyAttributeValue(element, "AXSize", None)
        if err == 0 and size_val is not None:
            ok, size = AS.AXValueGetValue(size_val, AS.kAXValueTypeCGSize, None)
            if ok:
                entry["width"] = round(size.width, 1)
                entry["height"] = round(size.height, 1)

        if entry:
            items.append(entry)

        # Children
        err, children = AS.AXUIElementCopyAttributeValue(element, "AXChildren", None)
        if err == 0 and children:
            for child in children:
                items.extend(_walk(child, depth + 1))

        return items

    return _walk(app_ref, 0)


# ---------------------------------------------------------------------------
# Mouse Control tools
# ---------------------------------------------------------------------------


@mcp.tool()
def mouse_click(
    x: float,
    y: float,
    button: str = "left",
    click_type: str = "single",
) -> str:
    """Click at screen coordinates (x, y).

    button: "left" or "right"
    click_type: "single" or "double"
    """
    point = Quartz.CGPointMake(x, y)

    if button == "right":
        down_type = Quartz.kCGEventRightMouseDown
        up_type = Quartz.kCGEventRightMouseUp
        mouse_button = Quartz.kCGMouseButtonRight
    else:
        down_type = Quartz.kCGEventLeftMouseDown
        up_type = Quartz.kCGEventLeftMouseUp
        mouse_button = Quartz.kCGMouseButtonLeft

    click_count = 2 if click_type == "double" else 1

    for i in range(click_count):
        event_down = Quartz.CGEventCreateMouseEvent(None, down_type, point, mouse_button)
        if click_type == "double":
            Quartz.CGEventSetIntegerValueField(
                event_down, Quartz.kCGMouseEventClickState, i + 1
            )
        Quartz.CGEventPost(Quartz.kCGHIDEventTap, event_down)

        event_up = Quartz.CGEventCreateMouseEvent(None, up_type, point, mouse_button)
        if click_type == "double":
            Quartz.CGEventSetIntegerValueField(
                event_up, Quartz.kCGMouseEventClickState, i + 1
            )
        Quartz.CGEventPost(Quartz.kCGHIDEventTap, event_up)

        if i < click_count - 1:
            time.sleep(0.05)

    return f"Clicked {button} ({click_type}) at ({x}, {y})"


@mcp.tool()
def mouse_move(x: float, y: float) -> str:
    """Move the mouse cursor to screen coordinates (x, y)."""
    point = Quartz.CGPointMake(x, y)
    event = Quartz.CGEventCreateMouseEvent(
        None, Quartz.kCGEventMouseMoved, point, Quartz.kCGMouseButtonLeft
    )
    Quartz.CGEventPost(Quartz.kCGHIDEventTap, event)
    return f"Moved cursor to ({x}, {y})"


@mcp.tool()
def mouse_drag(
    x1: float,
    y1: float,
    x2: float,
    y2: float,
    steps: int = 20,
    duration: float = 0.5,
) -> str:
    """Drag from (x1, y1) to (x2, y2) with smooth interpolation.

    steps: number of intermediate points (more = smoother, default 20)
    duration: total drag time in seconds (default 0.5)

    Great for drawing in apps like Preview, Figma, Photoshop, etc.
    """
    start = Quartz.CGPointMake(x1, y1)
    delay = duration / max(steps, 1)

    # Mouse down at start
    event = Quartz.CGEventCreateMouseEvent(
        None, Quartz.kCGEventLeftMouseDown, start, Quartz.kCGMouseButtonLeft
    )
    Quartz.CGEventPost(Quartz.kCGHIDEventTap, event)
    time.sleep(delay)

    # Intermediate drag points
    for i in range(1, steps + 1):
        t = i / steps
        ix = x1 + (x2 - x1) * t
        iy = y1 + (y2 - y1) * t
        pt = Quartz.CGPointMake(ix, iy)
        event = Quartz.CGEventCreateMouseEvent(
            None, Quartz.kCGEventLeftMouseDragged, pt, Quartz.kCGMouseButtonLeft
        )
        Quartz.CGEventPost(Quartz.kCGHIDEventTap, event)
        time.sleep(delay)

    # Mouse up at end
    end = Quartz.CGPointMake(x2, y2)
    event = Quartz.CGEventCreateMouseEvent(
        None, Quartz.kCGEventLeftMouseUp, end, Quartz.kCGMouseButtonLeft
    )
    Quartz.CGEventPost(Quartz.kCGHIDEventTap, event)

    return f"Dragged from ({x1}, {y1}) to ({x2}, {y2}) in {steps} steps"


@mcp.tool()
def scroll(
    x: float,
    y: float,
    dy: int = 0,
    dx: int = 0,
) -> str:
    """Scroll at screen coordinates (x, y).

    dy: vertical scroll amount (positive = up, negative = down)
    dx: horizontal scroll amount (positive = left, negative = right)
    """
    # Move cursor to scroll position first
    move_event = Quartz.CGEventCreateMouseEvent(
        None, Quartz.kCGEventMouseMoved, Quartz.CGPointMake(x, y), Quartz.kCGMouseButtonLeft
    )
    Quartz.CGEventPost(Quartz.kCGHIDEventTap, move_event)
    time.sleep(0.05)

    scroll_event = Quartz.CGEventCreateScrollWheelEvent(
        None, Quartz.kCGScrollEventUnitLine, 2, dy, dx
    )
    Quartz.CGEventPost(Quartz.kCGHIDEventTap, scroll_event)

    return f"Scrolled at ({x}, {y}) dy={dy} dx={dx}"


# ---------------------------------------------------------------------------
# Keyboard Control tools
# ---------------------------------------------------------------------------


def _char_to_keycode(ch: str) -> tuple[int, bool]:
    """Return (keycode, needs_shift) for a character."""
    lower = ch.lower()
    if lower in KEY_CODES:
        return KEY_CODES[lower], ch.isupper() or ch != lower
    if ch in SHIFT_CHARS:
        base = SHIFT_CHARS[ch]
        return KEY_CODES[base], True
    if ch == " ":
        return KEY_CODES["space"], False
    if ch == "\n":
        return KEY_CODES["return"], False
    if ch == "\t":
        return KEY_CODES["tab"], False
    raise ValueError(f"No key code mapping for character: {ch!r}")


@mcp.tool()
def type_text(text: str, delay: float = 0.02) -> str:
    """Type a string of text character by character using CGEvents.

    delay: seconds between keystrokes (default 0.02)
    """
    for ch in text:
        keycode, needs_shift = _char_to_keycode(ch)
        flags = Quartz.kCGEventFlagMaskShift if needs_shift else 0

        event_down = Quartz.CGEventCreateKeyboardEvent(None, keycode, True)
        if flags:
            Quartz.CGEventSetFlags(event_down, flags)
        Quartz.CGEventPost(Quartz.kCGHIDEventTap, event_down)

        event_up = Quartz.CGEventCreateKeyboardEvent(None, keycode, False)
        if flags:
            Quartz.CGEventSetFlags(event_up, flags)
        Quartz.CGEventPost(Quartz.kCGHIDEventTap, event_up)

        time.sleep(delay)

    return f"Typed {len(text)} characters"


@mcp.tool()
def key_press(keys: str) -> str:
    """Press a key combination like "cmd+c", "cmd+shift+z", "enter", "tab", etc.

    Supports modifiers: cmd/command, shift, alt/option/opt, ctrl/control, fn
    Supports keys: a-z, 0-9, f1-f12, enter/return, tab, space, escape/esc,
                   up/down/left/right, delete/backspace, home, end, pageup, pagedown,
                   and common punctuation.
    """
    parts = [p.strip().lower() for p in keys.split("+")]

    # Separate modifiers from the key
    modifier_mask = 0
    key_name = None
    for part in parts:
        if part in MODIFIER_FLAGS:
            modifier_mask |= MODIFIER_FLAGS[part]
        else:
            key_name = part

    if key_name is None:
        raise ValueError(f"No key found in combination: {keys!r}")

    if key_name not in KEY_CODES:
        raise ValueError(
            f"Unknown key: {key_name!r}. Available: {', '.join(sorted(KEY_CODES))}"
        )

    keycode = KEY_CODES[key_name]

    event_down = Quartz.CGEventCreateKeyboardEvent(None, keycode, True)
    if modifier_mask:
        Quartz.CGEventSetFlags(event_down, modifier_mask)
    Quartz.CGEventPost(Quartz.kCGHIDEventTap, event_down)

    event_up = Quartz.CGEventCreateKeyboardEvent(None, keycode, False)
    if modifier_mask:
        Quartz.CGEventSetFlags(event_up, modifier_mask)
    Quartz.CGEventPost(Quartz.kCGHIDEventTap, event_up)

    return f"Pressed {keys}"


# ---------------------------------------------------------------------------
# App Management tools
# ---------------------------------------------------------------------------


@mcp.tool()
def open_app(name: str | None = None, bundle_id: str | None = None) -> str:
    """Launch an application by name or bundle ID.

    Provide either name (e.g. "Safari") or bundle_id (e.g. "com.apple.Safari").
    """
    workspace = AppKit.NSWorkspace.sharedWorkspace()

    if bundle_id:
        url = workspace.URLForApplicationWithBundleIdentifier_(bundle_id)
        if url is None:
            raise ValueError(f"No app found with bundle ID: {bundle_id}")
        config = AppKit.NSWorkspaceOpenConfiguration.configuration()
        # Use synchronous-style: open and wait briefly
        workspace.openApplicationAtURL_configuration_completionHandler_(
            url, config, None
        )
        return f"Opened app with bundle ID: {bundle_id}"
    elif name:
        result = workspace.launchApplication_(name)
        if not result:
            raise RuntimeError(f"Failed to launch application: {name}")
        return f"Opened app: {name}"
    else:
        raise ValueError("Provide either name or bundle_id")


@mcp.tool()
def activate_app(
    name: str | None = None,
    bundle_id: str | None = None,
    pid: int | None = None,
) -> str:
    """Bring an application to the foreground.

    Identify the app by name, bundle_id, or pid.
    """
    workspace = AppKit.NSWorkspace.sharedWorkspace()
    apps = workspace.runningApplications()

    target = None
    for app in apps:
        if pid is not None and app.processIdentifier() == pid:
            target = app
            break
        if bundle_id and str(app.bundleIdentifier() or "") == bundle_id:
            target = app
            break
        if name and str(app.localizedName() or "").lower() == name.lower():
            target = app
            break

    if target is None:
        identifier = pid or bundle_id or name
        raise ValueError(f"No running app found matching: {identifier}")

    target.activateWithOptions_(AppKit.NSApplicationActivateIgnoringOtherApps)
    return f"Activated: {target.localizedName()} (PID {target.processIdentifier()})"


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------


def main():
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
