server.py

#!/usr/bin/env python3
"""Windows OS MCP Server - screen capture, mouse control, and keyboard input."""

import ctypes
import ctypes.wintypes
import io
import time
from typing import Optional

import mss
import pyautogui
import win32con
import win32gui
import win32ui
from mcp.server.fastmcp import FastMCP, Image as MCPImage
from PIL import Image as PILImage, ImageDraw, ImageFont

pyautogui.PAUSE = 0.05
pyautogui.FAILSAFE = True  # move mouse to top-left corner to abort

mcp = FastMCP("WindowsOSMCP")


# ---------------------------------------------------------------------------
# Internal helpers — screen capture
# ---------------------------------------------------------------------------

def _capture_region(left: int, top: int, width: int, height: int) -> bytes:
    with mss.MSS() as sct:
        region = {"top": top, "left": left, "width": width, "height": height}
        raw = sct.grab(region)
        img = PILImage.frombytes("RGB", raw.size, raw.bgra, "raw", "BGRX")
    buf = io.BytesIO()
    img.save(buf, format="PNG", optimize=True)
    return buf.getvalue()


def _capture_window_printwindow(hwnd: int) -> tuple[bytes, int, int]:
    """Capture a window by HWND using the PrintWindow API.

    Returns (png_bytes, width, height).  The bitmap dimensions match
    GetWindowRect (logical pixels), so image pixel (px, py) maps directly
    to screen logical coordinate (rect.left + px, rect.top + py) — the same
    space that pyautogui / mouse_click use for clicks.

    Works even when the window is partially off-screen or behind other windows.
    Falls back gracefully: if PrintWindow returns nothing useful the caller
    should retry with _capture_region().
    """
    rect = win32gui.GetWindowRect(hwnd)
    w = rect[2] - rect[0]
    h = rect[3] - rect[1]

    hwnd_dc = win32gui.GetWindowDC(hwnd)
    mfc_dc = win32ui.CreateDCFromHandle(hwnd_dc)
    compat_dc = mfc_dc.CreateCompatibleDC()
    bmp = win32ui.CreateBitmap()
    bmp.CreateCompatibleBitmap(mfc_dc, w, h)
    compat_dc.SelectObject(bmp)

    # PW_RENDERFULLCONTENT (2) — captures layered / GPU-accelerated content.
    ctypes.windll.user32.PrintWindow(hwnd, compat_dc.GetSafeHdc(), 2)

    bmpinfo = bmp.GetInfo()
    bmpstr = bmp.GetBitmapBits(True)
    img = PILImage.frombuffer(
        "RGB",
        (bmpinfo["bmWidth"], bmpinfo["bmHeight"]),
        bmpstr, "raw", "BGRX", 0, 1,
    )

    win32gui.DeleteObject(bmp.GetHandle())
    compat_dc.DeleteDC()
    mfc_dc.DeleteDC()
    win32gui.ReleaseDC(hwnd, hwnd_dc)

    buf = io.BytesIO()
    img.save(buf, format="PNG", optimize=True)
    return buf.getvalue(), bmpinfo["bmWidth"], bmpinfo["bmHeight"]


def _find_window(title: str) -> Optional[int]:
    """Return the first visible hwnd whose title contains *title* (case-insensitive)."""
    matches: list[int] = []
    needle = title.lower()

    def _cb(hwnd, _):
        if win32gui.IsWindowVisible(hwnd):
            if needle in win32gui.GetWindowText(hwnd).lower():
                matches.append(hwnd)

    win32gui.EnumWindows(_cb, None)
    return matches[0] if matches else None


# ---------------------------------------------------------------------------
# Internal helpers — multi-monitor
# ---------------------------------------------------------------------------

class _MONITORINFOEX(ctypes.Structure):
    _fields_ = [
        ("cbSize",    ctypes.c_ulong),
        ("rcMonitor", ctypes.wintypes.RECT),
        ("rcWork",    ctypes.wintypes.RECT),
        ("dwFlags",   ctypes.c_ulong),
        ("szDevice",  ctypes.c_wchar * 32),
    ]

_MONITORENUMPROC = ctypes.WINFUNCTYPE(
    ctypes.c_bool,
    ctypes.wintypes.HMONITOR,
    ctypes.wintypes.HDC,
    ctypes.POINTER(ctypes.wintypes.RECT),
    ctypes.wintypes.LPARAM,
)


def _get_monitors() -> list[dict]:
    """
    Enumerate all connected monitors.

    Uses mss as the canonical source for index order and physical geometry,
    then enriches each entry with DPI/scale data from Win32 GetDpiForMonitor.

    Returns a list where index 1 = primary monitor, matching mss.monitors[1..N].
    Each dict has:
      index         – 1-based, matches mss.monitors[] indices
      device        – Windows device string, e.g. '\\\\.\\DISPLAY1'
      name          – friendly name from mss (if available)
      is_primary    – bool
      dpi           – effective DPI (96=100%, 120=125%, 144=150%, 192=200%)
      scale_factor  – dpi / 96  (1.0, 1.25, 1.5, 2.0 …)
      left/top      – monitor offset in virtual desktop (pyautogui / mouse_click space)
      width/height  – logical dimensions  (same coordinate space as mouse_click)
      phys_width/phys_height – physical pixel dimensions (raw screenshot image size)
    """
    # Collect Win32 per-monitor data keyed by (left, top) position
    win32_by_pos: dict[tuple[int, int], dict] = {}

    def _cb(hmon, hdc, _lprect, _param):
        info = _MONITORINFOEX()
        info.cbSize = ctypes.sizeof(_MONITORINFOEX)
        ctypes.windll.user32.GetMonitorInfoW(hmon, ctypes.byref(info))

        dpi_x = ctypes.c_uint(96)
        try:
            ctypes.windll.shcore.GetDpiForMonitor(
                hmon, 0, ctypes.byref(dpi_x), ctypes.byref(ctypes.c_uint())
            )
        except Exception:
            pass

        r = info.rcMonitor
        win32_by_pos[(r.left, r.top)] = {
            "device":     info.szDevice.strip(),
            "is_primary": bool(info.dwFlags & 1),
            "dpi":        dpi_x.value,
        }
        return True

    ctypes.windll.user32.EnumDisplayMonitors(
        None, None, _MONITORENUMPROC(_cb), 0
    )

    # mss provides the canonical ordering (index 1 = primary) and physical geometry
    with mss.MSS() as sct:
        mss_mons = sct.monitors[1:]  # skip [0] = virtual bounding box

    result: list[dict] = []
    for idx, mss_m in enumerate(mss_mons, start=1):
        pos = (mss_m["left"], mss_m["top"])
        w32 = win32_by_pos.get(pos, {})

        dpi = w32.get("dpi", 96)
        scale = round(dpi / 96.0, 4)
        phys_w = mss_m["width"]
        phys_h = mss_m["height"]
        log_w = round(phys_w / scale)
        log_h = round(phys_h / scale)

        result.append({
            "index":        idx,
            "device":       w32.get("device", ""),
            "name":         mss_m.get("name", ""),
            "is_primary":   w32.get("is_primary", idx == 1),
            "dpi":          dpi,
            "scale_factor": scale,
            "left":         mss_m["left"],
            "top":          mss_m["top"],
            "width":        log_w,
            "height":       log_h,
            "phys_width":   phys_w,
            "phys_height":  phys_h,
        })

    return result


def _get_monitor_by_index(index: int) -> Optional[dict]:
    for m in _get_monitors():
        if m["index"] == index:
            return m
    return None


def _capture_monitor_logical(monitor: dict) -> tuple[bytes, int, int]:
    """
    Capture the monitor at physical resolution via mss, then downscale to
    logical resolution when DPI > 100%.

    Returns (png_bytes, logical_width, logical_height).

    After downscaling, image pixel (px, py) maps directly to screen coordinate
    (monitor['left'] + px, monitor['top'] + py) — no further conversion needed.
    """
    with mss.MSS() as sct:
        mss_m = sct.monitors[monitor["index"]]
        raw = sct.grab(mss_m)

    img = PILImage.frombytes("RGB", raw.size, raw.bgra, "raw", "BGRX")

    scale = monitor["scale_factor"]
    log_w, log_h = monitor["width"], monitor["height"]
    if abs(scale - 1.0) > 0.02:          # only resize when DPI != 100%
        img = img.resize((log_w, log_h), PILImage.LANCZOS)

    buf = io.BytesIO()
    img.save(buf, format="PNG", optimize=True)
    return buf.getvalue(), log_w, log_h


def _coord_hint(monitor: dict) -> str:
    """Return a one-line coordinate-mapping note for a monitor screenshot."""
    L, T = monitor["left"], monitor["top"]
    scale = monitor["scale_factor"]
    if abs(scale - 1.0) <= 0.02:
        return (
            f"IMAGE->CLICK: screen_x = {L} + image_x,  screen_y = {T} + image_y  "
            f"(100% DPI, no scaling)."
        )
    return (
        f"IMAGE->CLICK: screen_x = {L} + image_x,  screen_y = {T} + image_y  "
        f"(screenshot pre-scaled {monitor['phys_width']}x{monitor['phys_height']} -> "
        f"{monitor['width']}x{monitor['height']} logical px at {scale}x DPI)."
    )


def _draw_grid(
    img: PILImage.Image,
    offset_x: int = 0,
    offset_y: int = 0,
    step: int = 100,
) -> PILImage.Image:
    """Overlay a labeled coordinate grid on a copy of *img*.

    Vertical lines + x-label at top, horizontal lines + y-label at left.
    Labels show actual screen coordinates (offset + image position), so the
    model can read the click target directly off the image regardless of how
    much the preview is shrunk by the client.
    """
    img = img.copy()
    draw = ImageDraw.Draw(img)
    w, h = img.size
    try:
        font = ImageFont.load_default(size=11)
    except TypeError:
        font = ImageFont.load_default()

    LINE = (160, 160, 160)   # light gray lines
    TEXT = (255, 255, 0)     # yellow labels
    BG   = (0, 0, 0)         # black label background

    for ix in range(0, w, step):
        draw.line([(ix, 0), (ix, h)], fill=LINE, width=1)
        lbl = str(offset_x + ix)
        bb = draw.textbbox((ix + 2, 2), lbl, font=font)
        draw.rectangle(bb, fill=BG)
        draw.text((ix + 2, 2), lbl, fill=TEXT, font=font)

    for iy in range(0, h, step):
        draw.line([(0, iy), (w, iy)], fill=LINE, width=1)
        lbl = str(offset_y + iy)
        bb = draw.textbbox((2, iy + 2), lbl, font=font)
        draw.rectangle(bb, fill=BG)
        draw.text((2, iy + 2), lbl, fill=TEXT, font=font)

    return img


def _png_from_image(img: PILImage.Image) -> bytes:
    buf = io.BytesIO()
    img.save(buf, format="PNG", optimize=True)
    return buf.getvalue()


def _is_blank(img: PILImage.Image) -> bool:
    """True if every pixel is identical — PrintWindow's failure mode for
    WinUI3 / DirectComposition apps (modern Task Manager, Windows Terminal,
    some UWP apps) that don't paint into a GDI-visible surface: it returns
    a flat black (or occasionally white) frame instead of raising."""
    extrema = img.convert("RGB").getextrema()
    return all(lo == hi for lo, hi in extrema)


def _capture_window(hwnd: int, rect: tuple[int, int, int, int]) -> tuple[bytes, int, int]:
    """Capture a window, preferring PrintWindow (works even when the window
    is occluded or behind others) but falling back to a screen-region grab
    of the composited desktop when PrintWindow produces a blank frame or
    raises — see _is_blank(). The region fallback only sees whatever is
    actually on top on screen, but that beats a guaranteed-useless black image.
    """
    left, top, right, bottom = rect
    width, height = right - left, bottom - top
    try:
        png, w, h = _capture_window_printwindow(hwnd)
        if not _is_blank(PILImage.open(io.BytesIO(png))):
            return png, w, h
    except Exception:
        pass
    return _capture_region(left, top, width, height), width, height


# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------

@mcp.tool()
def screenshot(window_title: str = "", monitor_index: int = 0, grid: bool = False) -> list:
    """
    Take a screenshot of a monitor or a specific window.

    The returned image is pre-scaled to LOGICAL resolution so that
    image pixel (px, py) maps directly to the screen coordinates used
    by mouse_click / mouse_move — no manual conversion needed.

    Args:
        window_title:  Partial window title (case-insensitive). When given,
                       captures that window and ignores monitor_index.
        monitor_index: Which monitor to capture (1-based index from
                       list_monitors). 0 = primary monitor (default).
        grid:          When True, overlays gridlines every 100 logical px with
                       labeled screen coordinates so click targets can be read
                       directly off the image regardless of preview scaling.

    Returns text metadata followed by the PNG image.
    The metadata always includes the exact IMAGE→CLICK coordinate formula.
    """
    if window_title:
        hwnd = _find_window(window_title)
        if not hwnd:
            return [f"Error: no visible window matching '{window_title}'."]
        rect = win32gui.GetWindowRect(hwnd)
        left, top, right, bottom = rect
        width, height = right - left, bottom - top
        if width <= 0 or height <= 0:
            return ["Error: window has zero size (may be minimised)."]
        actual_title = win32gui.GetWindowText(hwnd)
        png, img_w, img_h = _capture_window(hwnd, rect)
        if grid:
            img_obj = PILImage.open(io.BytesIO(png))
            png = _png_from_image(_draw_grid(img_obj, offset_x=left, offset_y=top))
        return [
            (
                f"Window '{actual_title}' ({img_w}x{img_h} px). "
                f"Window top-left: ({left}, {top}). "
                f"IMAGE→CLICK: screen_x = {left} + image_x,  "
                f"screen_y = {top} + image_y."
            ),
            MCPImage(data=png, format="png"),
        ]

    # Monitor screenshot
    if monitor_index == 0:
        monitors = _get_monitors()
        m = next((mon for mon in monitors if mon["is_primary"]), monitors[0])
    else:
        m = _get_monitor_by_index(monitor_index)
        if m is None:
            return [f"Error: no monitor with index {monitor_index}. "
                    f"Call list_monitors() to see available monitors."]

    png, img_w, img_h = _capture_monitor_logical(m)
    if grid:
        img_obj = PILImage.open(io.BytesIO(png))
        png = _png_from_image(_draw_grid(img_obj, offset_x=m["left"], offset_y=m["top"]))
    label = "primary" if m["is_primary"] else f"monitor {m['index']}"
    return [
        (
            f"Monitor {m['index']} ({label}) - {m['width']}x{m['height']} px. "
            f"Screen area: left={m['left']}, top={m['top']}, "
            f"right={m['left'] + m['width']}, bottom={m['top'] + m['height']}. "
            f"{_coord_hint(m)}"
        ),
        MCPImage(data=png, format="png"),
    ]


@mcp.tool()
def list_monitors() -> list[dict]:
    """
    Return information about every connected monitor.

    Fields per monitor:
      index         – 1-based index; use with screenshot_monitor() and
                      image_to_screen_coords()
      device        – Windows device name (e.g. '\\\\.\\DISPLAY1')
      name          – friendly monitor name from the OS
      is_primary    – whether this is the primary/main display
      dpi           – effective DPI  (96=100%, 120=125%, 144=150%, 192=200%)
      scale_factor  – DPI / 96  (1.0 = no scaling, 1.5 = 150%, 2.0 = 200%)
      left / top    – top-left corner in the virtual desktop  ← same coordinate
      width / height – logical pixel dimensions                  space as mouse_click
      phys_width / phys_height – raw physical pixel dimensions (screenshot size
                                 before logical downscaling)

    IMPORTANT: 'left', 'top', 'width', 'height' are in the SAME coordinate space
    as mouse_click / mouse_move.  A monitor with left=-3840 means its left edge is
    at screen x = -3840 and you must use negative x-values to click on it.
    """
    monitors = _get_monitors()
    return [{k: v for k, v in m.items() if not k.startswith("_")} for m in monitors]


@mcp.tool()
def screenshot_monitor(monitor_index: int, grid: bool = False) -> list:
    """
    Take a screenshot of a specific monitor by its 1-based index.

    Use list_monitors() first to discover available monitor indices.

    The image is pre-scaled to logical resolution. To click at any visible
    position, simply use:
        mouse_click(monitor.left + image_x, monitor.top + image_y)

    Args:
        monitor_index: 1-based monitor index from list_monitors().
        grid:          When True, overlays gridlines every 100 logical px with
                       labeled screen coordinates.
    """
    m = _get_monitor_by_index(monitor_index)
    if m is None:
        return [f"Error: no monitor with index {monitor_index}. "
                f"Call list_monitors() to see available monitors."]

    png, img_w, img_h = _capture_monitor_logical(m)
    if grid:
        img_obj = PILImage.open(io.BytesIO(png))
        png = _png_from_image(_draw_grid(img_obj, offset_x=m["left"], offset_y=m["top"]))
    label = "PRIMARY" if m["is_primary"] else f"monitor {m['index']}"
    return [
        (
            f"[{label}] Monitor {m['index']} - {m.get('name') or m['device']} - "
            f"{m['dpi']} DPI ({m['scale_factor']}x scale). "
            f"Image: {img_w}x{img_h} px (logical). "
            f"Screen area: left={m['left']}, top={m['top']}, "
            f"right={m['left'] + m['width']}, bottom={m['top'] + m['height']}. "
            f"{_coord_hint(m)}"
        ),
        MCPImage(data=png, format="png"),
    ]


@mcp.tool()
def screenshot_all_monitors() -> list:
    """
    Take a screenshot of every connected monitor simultaneously.

    Returns interleaved text metadata and PNG images — one text+image pair
    per monitor.  Each image is pre-scaled to logical resolution so that
    image pixel (px, py) maps directly to (monitor.left + px, monitor.top + py)
    for mouse_click — no coordinate conversion required.

    Use this to get a complete view of the entire desktop across all screens.
    """
    monitors = _get_monitors()
    result = []
    for m in monitors:
        png, img_w, img_h = _capture_monitor_logical(m)
        label = "PRIMARY" if m["is_primary"] else f"monitor {m['index']}"
        result.append(
            f"[{label}] Monitor {m['index']} - {m.get('name') or m['device']} - "
            f"{m['dpi']} DPI ({m['scale_factor']}x). "
            f"Image: {img_w}x{img_h} px. "
            f"Screen area: left={m['left']}, top={m['top']}, "
            f"right={m['left'] + m['width']}, bottom={m['top'] + m['height']}. "
            f"{_coord_hint(m)}"
        )
        result.append(MCPImage(data=png, format="png"))
    return result


@mcp.tool()
def screenshot_region(
    left: int,
    top: int,
    width: int,
    height: int,
    grid: bool = False,
) -> list:
    """
    Capture a rectangular sub-region of the screen at full physical resolution.

    Unlike screenshot(), no downscaling is applied — on HiDPI monitors you get
    the raw physical pixels, which preserves fine detail that would otherwise
    be lost when a full-screen capture is shrunk to fit the preview.

    Coordinates (left, top, width, height) are in logical pixels — the same
    space as mouse_click.  On a 2x DPI monitor a 200x100 logical region yields
    a ~400x200 image; on 100% DPI the image is 200x100.

    IMAGE→CLICK conversion:
      screen_x = left + image_x / scale_factor
      screen_y = top  + image_y / scale_factor
    The metadata includes scale_factor so you can compute exact click coords.
    Use image_to_screen_coords() when you need precise conversions.

    Args:
        left, top:     Top-left corner in screen (logical) coordinates.
        width, height: Region size in logical pixels.
        grid:          Overlay a labeled coordinate grid (every 100 px,
                       labels show screen coordinates).
    """
    if width <= 0 or height <= 0:
        return ["Error: width and height must be positive."]

    monitors = _get_monitors()
    cx, cy = left + width // 2, top + height // 2
    mon = next(
        (m for m in monitors
         if m["left"] <= cx < m["left"] + m["width"]
         and m["top"] <= cy < m["top"] + m["height"]),
        next((m for m in monitors if m["is_primary"]), monitors[0]),
    )
    scale = mon["scale_factor"]

    png = _capture_region(left, top, width, height)
    img = PILImage.open(io.BytesIO(png))
    phys_w, phys_h = img.size

    if grid:
        img = _draw_grid(img, offset_x=left, offset_y=top)
        png = _png_from_image(img)

    if abs(scale - 1.0) > 0.02 and phys_w > width:
        click_hint = (
            f"IMAGE→CLICK: screen_x = {left} + image_x / {scale},  "
            f"screen_y = {top} + image_y / {scale}  "
            f"(physical {phys_w}x{phys_h} px at {scale}x DPI)."
        )
    else:
        click_hint = (
            f"IMAGE→CLICK: screen_x = {left} + image_x,  "
            f"screen_y = {top} + image_y."
        )

    return [
        (
            f"Region ({left},{top}) size {width}x{height} logical px. "
            f"Image: {phys_w}x{phys_h} px. "
            f"Monitor {mon['index']} ({mon['scale_factor']}x DPI). "
            f"{click_hint}"
        ),
        MCPImage(data=png, format="png"),
    ]


@mcp.tool()
def image_to_screen_coords(
    image_x: int,
    image_y: int,
    monitor_index: int,
) -> dict:
    """
    Convert image pixel coordinates from a monitor screenshot to screen
    coordinates suitable for mouse_click / mouse_move.

    Use this when you have (image_x, image_y) from a screenshot taken with
    screenshot_monitor() or screenshot() and want the exact screen position.

    Since all monitor screenshots are pre-scaled to logical resolution the
    formula is: screen_x = monitor.left + image_x,
                screen_y = monitor.top  + image_y.
    This tool makes the conversion explicit and catches out-of-range errors.

    Args:
        image_x:       X coordinate in the screenshot image (0 = leftmost column).
        image_y:       Y coordinate in the screenshot image (0 = top row).
        monitor_index: 1-based index of the monitor the screenshot came from
                       (from list_monitors()).
    """
    m = _get_monitor_by_index(monitor_index)
    if m is None:
        return {"error": f"No monitor with index {monitor_index}. "
                         "Call list_monitors() to see available monitors."}

    if not (0 <= image_x < m["width"]):
        return {"error": f"image_x={image_x} is out of range [0, {m['width']})."}
    if not (0 <= image_y < m["height"]):
        return {"error": f"image_y={image_y} is out of range [0, {m['height']})."}

    screen_x = m["left"] + image_x
    screen_y = m["top"]  + image_y
    return {
        "screen_x":       screen_x,
        "screen_y":       screen_y,
        "monitor_index":  monitor_index,
        "monitor_offset": {"left": m["left"], "top": m["top"]},
        "usage":          f"mouse_click({screen_x}, {screen_y})",
    }


@mcp.tool()
def list_windows() -> list[dict]:
    """
    Return all visible windows that have a non-empty title.

    Each entry contains: title, hwnd, rect (left/top/right/bottom/width/height),
    and monitor_index indicating which monitor the window centre falls on.
    """
    monitors = _get_monitors()

    def _monitor_for_point(cx: int, cy: int) -> int:
        for m in monitors:
            if m["left"] <= cx < m["left"] + m["width"] and \
               m["top"]  <= cy < m["top"]  + m["height"]:
                return m["index"]
        return 1  # fallback to primary

    windows: list[dict] = []

    def _cb(hwnd, _):
        if win32gui.IsWindowVisible(hwnd):
            title = win32gui.GetWindowText(hwnd).strip()
            if title:
                r = win32gui.GetWindowRect(hwnd)
                cx = (r[0] + r[2]) // 2
                cy = (r[1] + r[3]) // 2
                windows.append({
                    "title": title,
                    "hwnd": hwnd,
                    "monitor_index": _monitor_for_point(cx, cy),
                    "rect": {
                        "left": r[0], "top": r[1],
                        "right": r[2], "bottom": r[3],
                        "width": r[2] - r[0],
                        "height": r[3] - r[1],
                    },
                })

    win32gui.EnumWindows(_cb, None)
    return sorted(windows, key=lambda w: w["title"].lower())


@mcp.tool()
def focus_window(title: str) -> dict:
    """
    Bring a window to the foreground by partial title match.

    Args:
        title: Partial window title (case-insensitive).
    """
    hwnd = _find_window(title)
    if not hwnd:
        return {"success": False, "error": f"No window matching '{title}'."}
    try:
        placement = win32gui.GetWindowPlacement(hwnd)
        if placement[1] == win32con.SW_SHOWMINIMIZED:
            win32gui.ShowWindow(hwnd, win32con.SW_RESTORE)
        win32gui.SetForegroundWindow(hwnd)
        time.sleep(0.15)
        return {"success": True, "window": win32gui.GetWindowText(hwnd), "hwnd": hwnd}
    except Exception as exc:
        return {"success": False, "error": str(exc)}


@mcp.tool()
def get_window_rect(title: str) -> dict:
    """
    Return the position and size of a window by partial title match.

    Includes convenience fields: center (screen coords of the middle) and
    monitor_index indicating which monitor the window is primarily on.
    """
    hwnd = _find_window(title)
    if not hwnd:
        return {"error": f"No window matching '{title}'."}
    r = win32gui.GetWindowRect(hwnd)
    cx = (r[0] + r[2]) // 2
    cy = (r[1] + r[3]) // 2

    monitors = _get_monitors()
    mon_idx = 1
    for m in monitors:
        if m["left"] <= cx < m["left"] + m["width"] and \
           m["top"]  <= cy < m["top"]  + m["height"]:
            mon_idx = m["index"]
            break

    return {
        "title": win32gui.GetWindowText(hwnd),
        "left": r[0], "top": r[1], "right": r[2], "bottom": r[3],
        "width": r[2] - r[0],
        "height": r[3] - r[1],
        "center": {"x": cx, "y": cy},
        "monitor_index": mon_idx,
    }


@mcp.tool()
def mouse_move(x: int, y: int, duration: float = 0.2) -> dict:
    """
    Move the mouse cursor to absolute screen coordinates.

    Args:
        x: Horizontal screen coordinate in pixels.
        y: Vertical screen coordinate in pixels.
        duration: Seconds the move animation takes (0 = instant).
    """
    pyautogui.moveTo(x, y, duration=duration)
    pos = pyautogui.position()
    return {"position": {"x": pos.x, "y": pos.y}}


@mcp.tool()
def mouse_click(
    x: int,
    y: int,
    button: str = "left",
    double: bool = False,
    clicks: int = 1,
) -> dict:
    """
    Click the mouse at screen coordinates.

    Args:
        x, y: Screen coordinates.
        button: 'left', 'right', or 'middle'.
        double: True to double-click (overrides clicks).
        clicks: Number of clicks (ignored when double=True).
    """
    if button not in ("left", "right", "middle"):
        return {"error": f"Invalid button '{button}'. Use left/right/middle."}
    if double:
        pyautogui.doubleClick(x, y, button=button)
        action = "double_click"
    else:
        pyautogui.click(x, y, button=button, clicks=clicks)
        action = "click" if clicks == 1 else f"{clicks}x_click"
    return {"action": action, "button": button, "x": x, "y": y}


@mcp.tool()
def mouse_drag(
    from_x: int,
    from_y: int,
    to_x: int,
    to_y: int,
    button: str = "left",
    duration: float = 0.4,
) -> dict:
    """
    Press-hold-drag from one screen position to another, then release.

    Args:
        from_x, from_y: Starting coordinates.
        to_x, to_y: Ending coordinates.
        button: Mouse button held during drag.
        duration: Seconds the drag takes.
    """
    pyautogui.moveTo(from_x, from_y, duration=0.1)
    pyautogui.dragTo(to_x, to_y, duration=duration, button=button)
    return {"from": {"x": from_x, "y": from_y}, "to": {"x": to_x, "y": to_y}}


@mcp.tool()
def mouse_scroll(x: int, y: int, clicks: int, direction: str = "vertical") -> dict:
    """
    Scroll the mouse wheel at the given coordinates.

    Args:
        x, y:      Screen coordinates to scroll at.
        clicks:    Scroll amount — positive scrolls up/right, negative scrolls down/left.
        direction: 'vertical' (default) or 'horizontal'.
    """
    if direction not in ("vertical", "horizontal"):
        return {"error": f"Invalid direction '{direction}'. Use 'vertical' or 'horizontal'."}
    pyautogui.moveTo(x, y, duration=0.1)
    if direction == "horizontal":
        pyautogui.hscroll(clicks)
    else:
        pyautogui.scroll(clicks)
    return {"scrolled": clicks, "direction": direction, "x": x, "y": y}


@mcp.tool()
def keyboard_type(text: str, interval: float = 0.04) -> dict:
    """
    Type a string of text into the currently focused window.

    For arbitrary Unicode, the text is pushed via the clipboard so all characters
    are supported. The original clipboard contents are restored afterwards.

    Args:
        text: The text to type.
        interval: Seconds between keystrokes (only used for ASCII fallback).
    """
    import subprocess
    # Use PowerShell to set clipboard so we don't need pyperclip
    escaped = text.replace("'", "''")
    subprocess.run(
        ["powershell", "-NoProfile", "-Command", f"Set-Clipboard -Value '{escaped}'"],
        check=True,
        capture_output=True,
    )
    time.sleep(0.05)
    pyautogui.hotkey("ctrl", "v")
    return {"typed": text, "length": len(text), "method": "clipboard_paste"}


@mcp.tool()
def keyboard_type_ascii(text: str, interval: float = 0.04) -> dict:
    """
    Type ASCII text character-by-character using pyautogui keystrokes.

    Prefer keyboard_type for Unicode. Use this when you need per-keystroke
    timing control (e.g. typing into a terminal that processes each key).

    Args:
        text: ASCII text to type.
        interval: Seconds between individual keystrokes.
    """
    pyautogui.write(text, interval=interval)
    return {"typed": text, "length": len(text), "method": "keystroke"}


@mcp.tool()
def key_press(keys: str) -> dict:
    """
    Press a key or keyboard shortcut.

    Single key examples: 'enter', 'escape', 'tab', 'space', 'backspace',
                         'delete', 'up', 'down', 'left', 'right', 'f5'
    Combo examples: 'ctrl+c', 'ctrl+shift+t', 'alt+f4', 'win+d', 'ctrl+alt+delete'

    Separate keys with '+'. Call multiple times for key sequences.

    Args:
        keys: Key name or '+'-separated combo string.
    """
    parts = [k.strip().lower() for k in keys.split("+")]
    if len(parts) == 1:
        pyautogui.press(parts[0])
    else:
        pyautogui.hotkey(*parts)
    return {"pressed": keys, "parts": parts}


@mcp.tool()
def get_cursor_position() -> dict:
    """Return the current mouse cursor position in screen coordinates."""
    pos = pyautogui.position()
    return {"x": pos.x, "y": pos.y}


# ---------------------------------------------------------------------------
# Internal helpers — UI Automation
# ---------------------------------------------------------------------------

# ControlTypes worth numbering in screenshot_elements() by default — the
# things a user actually clicks/types into, as opposed to layout containers
# (Pane, Group, Window, Custom, ScrollBar, TitleBar, ...) which just add
# visual noise to an annotated screenshot.
_INTERESTING_CONTROL_TYPES = {
    "button", "edit", "checkbox", "combobox", "list", "listitem",
    "tree", "treeitem", "menu", "menuitem", "tabitem", "radiobutton",
    "slider", "spinner", "hyperlink", "splitbutton", "dataitem",
    "headeritem", "progressbar",
}


def _get_uia_descendants(hwnd: int) -> list:
    """Connect to *hwnd* via UIA and return all descendant elements.

    Raises ImportError if pywinauto isn't installed, or any other exception
    if the UIA connection fails — callers should catch both.
    """
    from pywinauto import Application  # lazy import — not needed for other tools
    app = Application(backend="uia").connect(handle=hwnd)
    win = app.window(handle=hwnd)
    return win.descendants()


def _elem_summary(elem) -> Optional[dict]:
    """Extract a JSON-safe summary from a pywinauto UIA element, or None if
    the element has no meaningful (non-zero-size) rect or can't be read."""
    try:
        rect = elem.rectangle()
        if rect.width() <= 0 or rect.height() <= 0:
            return None
        return {
            "name":           elem.window_text() or "",
            "automation_id":  elem.automation_id() or "",
            "control_type":   (elem.element_info.control_type or "").strip(),
            "enabled":        elem.is_enabled(),
            "visible":        elem.is_visible(),
            "rect": {
                "left":     rect.left,
                "top":      rect.top,
                "right":    rect.right,
                "bottom":   rect.bottom,
                "center_x": (rect.left + rect.right)  // 2,
                "center_y": (rect.top  + rect.bottom) // 2,
            },
        }
    except Exception:
        return None


@mcp.tool()
def find_element(
    window_title: str,
    automation_id: str = "",
    name: str = "",
    control_type: str = "",
) -> list[dict]:
    """
    Find UI elements inside a window using UI Automation (UIA).

    Targeting by automation_id or name is immune to any pixel / DPI issues
    because it queries the accessibility tree, not the screen image.
    Use this instead of eyeballing screenshots whenever the app exposes UIA
    (WinForms, WPF, Win32 with accessibility, most native Windows apps do).

    Args:
        window_title:  Partial window title (case-insensitive).
        automation_id: Exact AutomationId to match (empty = any).
        name:          Partial element Name/title to match (case-insensitive,
                       empty = any).
        control_type:  ControlType to filter by, e.g. 'Button', 'Edit',
                       'Text', 'CheckBox', 'ComboBox', 'List', 'ListItem',
                       'Tree', 'TreeItem', 'Menu', 'MenuItem', 'ToolBar',
                       'TabControl', 'TabItem', 'GroupBox', 'RadioButton',
                       'Slider', 'Spinner', 'Window', 'Pane', 'Document'.
                       Empty = any.

    Returns a list of matching elements.  Each element includes its screen
    rect (center_x / center_y are ready to pass to mouse_click).
    """
    hwnd = _find_window(window_title)
    if not hwnd:
        return [{"error": f"No visible window matching '{window_title}'."}]

    try:
        elements = _get_uia_descendants(hwnd)
    except ImportError:
        return [{"error": "pywinauto is not installed. Run: pip install pywinauto"}]
    except Exception as exc:
        return [{"error": f"UIA connection failed: {exc}"}]

    results: list[dict] = []
    for elem in elements:
        info = _elem_summary(elem)
        if info is None:
            continue
        if automation_id and info["automation_id"] != automation_id:
            continue
        if name          and name.lower() not in info["name"].lower():
            continue
        if control_type  and info["control_type"].lower() != control_type.lower():
            continue
        results.append(info)

    if not results:
        return [{"message": "No matching elements found.", "total_searched": len(elements)}]
    return results[:100]


@mcp.tool()
def click_element(
    window_title: str,
    automation_id: str = "",
    name: str = "",
    control_type: str = "",
) -> dict:
    """
    Click the first UI element matching the given criteria via UI Automation.

    This is the most reliable way to click controls — it targets by
    automation_id or name rather than pixel coordinates, so it works
    correctly at any DPI and even when the window is partially off-screen.

    Provide at least one of automation_id, name, or control_type to narrow
    the match.  If multiple elements match, the first one in the UIA tree is
    clicked.

    Args:
        window_title:  Partial window title (case-insensitive).
        automation_id: Exact AutomationId to match.
        name:          Partial element Name to match (case-insensitive).
        control_type:  ControlType string (e.g. 'Button', 'MenuItem').
    """
    hwnd = _find_window(window_title)
    if not hwnd:
        return {"success": False, "error": f"No visible window matching '{window_title}'."}

    try:
        elements = _get_uia_descendants(hwnd)
    except ImportError:
        return {"success": False, "error": "pywinauto is not installed. Run: pip install pywinauto"}
    except Exception as exc:
        return {"success": False, "error": f"UIA connection failed: {exc}"}

    for elem in elements:
        try:
            eid   = elem.automation_id() or ""
            ename = elem.window_text()    or ""
            etype = (elem.element_info.control_type or "").strip()

            if automation_id and eid   != automation_id:
                continue
            if name          and name.lower() not in ename.lower():
                continue
            if control_type  and etype.lower() != control_type.lower():
                continue

            rect = elem.rectangle()
            cx   = (rect.left + rect.right)  // 2
            cy   = (rect.top  + rect.bottom) // 2
            elem.click_input()
            return {
                "success":        True,
                "clicked_name":   ename,
                "automation_id":  eid,
                "control_type":   etype,
                "screen_x":       cx,
                "screen_y":       cy,
            }
        except Exception as exc:
            return {"success": False, "error": f"Click failed: {exc}"}

    return {"success": False, "error": "No matching element found."}


@mcp.tool()
def list_elements(
    window_title: str,
    control_type: str = "",
    visible_only: bool = True,
    enabled_only: bool = False,
    max_elements: int = 300,
) -> list[dict]:
    """
    Dump every control in a window's UI Automation tree as structured data.

    Unlike find_element (meant for targeted lookups, capped at 100 matches),
    this returns a broad inventory of the window's controls so you can
    reason about layout deterministically instead of guessing from a
    screenshot. Elements with a zero-size rect (not actually laid out) are
    always skipped.

    Args:
        window_title:  Partial window title (case-insensitive).
        control_type:  Optional exact ControlType filter (e.g. 'Button').
                       Empty = all types.
        visible_only:  Skip elements UIA reports as not visible (default True).
        enabled_only:  Skip disabled elements (default False).
        max_elements:  Safety cap on how many elements to return.

    Returns each element's name, automation_id, control_type, enabled/visible
    flags, and screen rect (center_x/center_y ready for mouse_click).
    """
    hwnd = _find_window(window_title)
    if not hwnd:
        return [{"error": f"No visible window matching '{window_title}'."}]

    try:
        elements = _get_uia_descendants(hwnd)
    except ImportError:
        return [{"error": "pywinauto is not installed. Run: pip install pywinauto"}]
    except Exception as exc:
        return [{"error": f"UIA connection failed: {exc}"}]

    results: list[dict] = []
    for elem in elements:
        info = _elem_summary(elem)
        if info is None:
            continue
        if control_type  and info["control_type"].lower() != control_type.lower():
            continue
        if visible_only   and not info["visible"]:
            continue
        if enabled_only   and not info["enabled"]:
            continue
        results.append(info)
        if len(results) >= max_elements:
            break

    if not results:
        return [{"message": "No matching elements found.", "total_searched": len(elements)}]
    return results


@mcp.tool()
def screenshot_elements(
    window_title: str,
    control_type: str = "",
    max_elements: int = 60,
) -> list:
    """
    Screenshot a window with numbered boxes drawn on every interactive UIA
    control, plus a legend mapping each number to its control info.

    This is the most reliable way to visually locate a control: instead of
    eyeballing pixels or guessing coordinates, read the number off the image
    and look up its exact click coordinates in the legend. Only works for
    apps that expose UIA well (WinForms, WPF, most native Windows apps) — for
    apps that don't, fall back to eyeballing a plain screenshot(grid=True).

    Args:
        window_title:  Partial window title (case-insensitive).
        control_type:  Optional exact ControlType filter to narrow which
                       controls get numbered (e.g. 'Button'). Empty = the
                       default set of interactive types (Button, Edit,
                       CheckBox, ComboBox, ListItem, TreeItem, MenuItem,
                       TabItem, RadioButton, Slider, Hyperlink, etc.) —
                       layout containers like Pane/Group are excluded to
                       keep the image readable.
        max_elements:  Cap on how many controls to number (keeps the image
                       readable). If more exist, use list_elements for the
                       full inventory or narrow with control_type.

    Returns text metadata + legend, followed by the annotated PNG image.
    """
    hwnd = _find_window(window_title)
    if not hwnd:
        return [f"Error: no visible window matching '{window_title}'."]

    try:
        elements = _get_uia_descendants(hwnd)
    except ImportError:
        return ["Error: pywinauto is not installed. Run: pip install pywinauto"]
    except Exception as exc:
        return [f"Error: UIA connection failed: {exc}"]

    rect = win32gui.GetWindowRect(hwnd)
    left, top, right, bottom = rect
    width, height = right - left, bottom - top
    if width <= 0 or height <= 0:
        return ["Error: window has zero size (may be minimised)."]
    actual_title = win32gui.GetWindowText(hwnd)
    png, img_w, img_h = _capture_window(hwnd, rect)

    candidates: list[dict] = []
    for elem in elements:
        info = _elem_summary(elem)
        if info is None or not info["visible"]:
            continue
        etype_l = info["control_type"].lower()
        if control_type:
            if etype_l != control_type.lower():
                continue
        elif etype_l not in _INTERESTING_CONTROL_TYPES:
            continue
        r = info["rect"]
        # keep only elements that actually fall within the captured window
        if r["right"] <= left or r["left"] >= right or r["bottom"] <= top or r["top"] >= bottom:
            continue
        candidates.append(info)

    # stable reading order: top-to-bottom, then left-to-right
    candidates.sort(key=lambda c: (c["rect"]["top"], c["rect"]["left"]))

    truncated = len(candidates) > max_elements
    candidates = candidates[:max_elements]

    img = PILImage.open(io.BytesIO(png))
    draw = ImageDraw.Draw(img)
    try:
        font = ImageFont.load_default(size=13)
    except TypeError:
        font = ImageFont.load_default()

    BOX  = (255, 0, 255)   # magenta box outline
    TEXT = (255, 255, 0)   # yellow label text
    BG   = (0, 0, 0)       # black label background

    legend_lines = []
    for i, info in enumerate(candidates, start=1):
        r = info["rect"]
        ix0, iy0 = r["left"] - left, r["top"] - top
        ix1, iy1 = r["right"] - left, r["bottom"] - top
        draw.rectangle([ix0, iy0, ix1, iy1], outline=BOX, width=2)
        label = str(i)
        bb = draw.textbbox((ix0 + 2, iy0 + 1), label, font=font)
        draw.rectangle(bb, fill=BG)
        draw.text((ix0 + 2, iy0 + 1), label, fill=TEXT, font=font)

        name_part = f" '{info['name']}'" if info["name"] else ""
        id_part   = f" id={info['automation_id']}" if info["automation_id"] else ""
        legend_lines.append(
            f"[{i}] {info['control_type']}{name_part}{id_part} "
            f"center=({r['center_x']},{r['center_y']})"
        )

    png = _png_from_image(img)

    meta = (
        f"Window '{actual_title}' ({img_w}x{img_h} px), top-left ({left},{top}). "
        f"{len(candidates)} controls numbered"
        + (" (truncated — more exist; narrow control_type or use list_elements for the rest)."
           if truncated else ".")
        + f" IMAGE→CLICK: screen_x = {left} + image_x,  screen_y = {top} + image_y.\n"
        + "Legend:\n" + "\n".join(legend_lines)
    )

    return [meta, MCPImage(data=png, format="png")]


if __name__ == "__main__":
    mcp.run()