"""Windows in-process desktop-control MCP for Jarvis.

This is the Windows replacement for the macOS ``macos-mcp`` HTTP server. It
gives Claude OS-wide "control any app" hands — see the screen and drive any
application (Teams, WhatsApp, Chrome, VS Code, anything) by moving the real
mouse, typing, pressing hotkeys, launching apps, listing windows, reading the
accessibility tree, and running shell commands.

Unlike the macOS build (which needed an external HTTP server so it could
inherit the app's Accessibility grant), Windows needs no permission wall for
SendInput/screenshots against ordinary non-elevated apps. So we run the tools
IN-PROCESS via the Claude Agent SDK's in-process MCP server: no external
server, no auth token, no HTTP hop. The server object is placed in the session
options ``mcp_servers`` dict under the key ``"desktop"``, and Claude sees the
tools as ``mcp__desktop__<name>``.

Backing libraries (all imported lazily inside each tool so importing this
module never hard-fails when a dependency is missing):
- ``mss``          — fast full-screen screenshots.
- ``pyautogui``    — mouse move/click, keyboard typing, hotkeys, scroll.
- ``pywin32``      — enumerate visible top-level windows and their bounds.
- ``uiautomation`` — read the foreground window's UI accessibility tree.

Set ``MCP_DESKTOP=0`` to disable the desktop tools entirely.
"""

from __future__ import annotations

import base64
import io
import os
import subprocess

import numpy as np


def enabled() -> bool:
    """Whether the desktop-control MCP should be exposed to Claude.
    Windows-only: the Mac uses the macos-mcp HTTP server instead
    (voice/desktop_server.py) — this module's hands are pywin32-based."""
    return os.name == "nt" and os.environ.get("MCP_DESKTOP", "1") != "0"


def shell_launch(name: str) -> tuple[bool, str]:
    """Open an app by fuzzy name the way the Start menu does.

    Resolves the name via `Get-StartApps` (the real Start-menu index, which
    covers Store/UWP apps, Steam games, and desktop apps alike) to an AppID,
    then launches it — a Steam/URL AppID directly, an AUMID via
    shell:AppsFolder. Falls back to Start-Process on the raw name. Returns
    (ok, matched-name-or-error). This is why "open Discord" / "open Marvel
    Rivals" works where a bare `start "Name"` fails.
    """
    ps = (
        '$n=$env:JARVIS_APP;'
        '$a=Get-StartApps | Where-Object { $_.Name -like "*$n*" } |'
        ' Sort-Object { $_.Name.Length } | Select-Object -First 1;'
        'if($a){ if($a.AppID -match "://"){ Start-Process $a.AppID }'
        ' else { Start-Process ("shell:AppsFolder\\" + $a.AppID) };'
        ' "OK:"+$a.Name }'
        ' else { try{ Start-Process $n; "OK:"+$n }catch{ "FAIL:"+$n } }'
    )
    try:
        env = dict(os.environ, JARVIS_APP=name)
        r = subprocess.run(
            ["powershell", "-NoProfile", "-Command", ps],
            env=env, capture_output=True, text=True, timeout=25,
        )
        out = (r.stdout or "").strip().splitlines()[-1] if r.stdout else ""
        return (out.startswith("OK:"), out.split(":", 1)[-1] or name)
    except Exception as e:  # noqa: BLE001
        return (False, str(e)[:80])


def _text(msg: str) -> dict:
    """Wrap a plain string as an MCP text content result."""
    return {"content": [{"type": "text", "text": msg}]}


# Downscale wide captures so an ultrawide screenshot doesn't blow the model's
# image budget — and ALWAYS tell the model the real size + scale factor so it
# can convert image coordinates back to true screen pixels before clicking.
_MAX_SHOT_W = 1568  # Anthropic's recommended max image edge


def _grab_screen() -> tuple[bytes, str]:
    """Capture the full virtual desktop → (png_bytes, geometry_note)."""
    import mss  # lazy
    with mss.mss() as sct:
        mon = sct.monitors[0]  # full virtual desktop (all monitors)
        shot = sct.grab(mon)
        w, h = shot.width, shot.height
        img = np.array(shot)[:, :, :3][:, :, ::-1]  # BGRA -> RGB
    scale = 1.0
    if w > _MAX_SHOT_W:
        scale = w / _MAX_SHOT_W
        import cv2  # lazy
        img = cv2.resize(img, (_MAX_SHOT_W, int(h / scale)),
                         interpolation=cv2.INTER_AREA)
    import cv2  # lazy
    ok, png = cv2.imencode(".png", img[:, :, ::-1])  # RGB -> BGR for encode
    if not ok:
        raise RuntimeError("png encode failed")
    note = (f"Real screen: {w}x{h} px (virtual desktop, origin "
            f"{mon['left']},{mon['top']}). Image shown: {img.shape[1]}x"
            f"{img.shape[0]} px — scale factor {scale:.3f}. To CLICK "
            f"something you located on this image, MULTIPLY its image "
            f"coordinates by {scale:.3f} to get real screen pixels."
            if scale > 1.0 else
            f"Screen: {w}x{h} px, shown at full resolution — image "
            f"coordinates ARE screen coordinates.")
    return png.tobytes(), note


def build_server():
    """Build and return the in-process ``desktop`` MCP server.

    Imports the SDK's in-process MCP helpers here (not at module top) so that
    merely importing this module never fails if the installed
    ``claude_agent_sdk`` lacks ``tool``/``create_sdk_mcp_server``. Returns the
    server object to place under ``mcp_servers["desktop"]``.
    """
    from claude_agent_sdk import tool, create_sdk_mcp_server

    @tool("screenshot", "Capture the full screen so you can SEE it. Returns "
          "the desktop image plus its exact scale factor for converting image "
          "coordinates to real screen pixels.", {})
    async def screenshot(args: dict) -> dict:
        try:
            png, note = _grab_screen()
            data = base64.b64encode(png).decode("ascii")
            return {"content": [
                {"type": "text", "text": note},
                {"type": "image", "data": data, "mimeType": "image/png"},
            ]}
        except Exception as e:  # noqa: BLE001
            return _text(f"screenshot failed: {e}")

    @tool("click", "Move the mouse to absolute screen pixel (x, y) and click. "
          "Origin is the top-left of the screen.",
          {"x": int, "y": int, "button": str, "double": bool})
    async def click(args: dict) -> dict:
        try:
            import pyautogui  # lazy
            x = int(args["x"])
            y = int(args["y"])
            button = args.get("button") or "left"
            double = bool(args.get("double", False))
            pyautogui.click(x, y, button=button, clicks=2 if double else 1)
            kind = "double-clicked" if double else "clicked"
            return _text(f"{kind} {button} at ({x}, {y}).")
        except Exception as e:  # noqa: BLE001
            return _text(f"click failed: {e}")

    @tool("move", "Move the mouse to absolute screen pixel (x, y).",
          {"x": int, "y": int})
    async def move(args: dict) -> dict:
        try:
            import pyautogui  # lazy
            x = int(args["x"])
            y = int(args["y"])
            pyautogui.moveTo(x, y)
            return _text(f"moved to ({x}, {y}).")
        except Exception as e:  # noqa: BLE001
            return _text(f"move failed: {e}")

    @tool("type_text", "Type the given text at the current focus, as if typed "
          "on the keyboard.", {"text": str})
    async def type_text(args: dict) -> dict:
        try:
            import pyautogui  # lazy
            text = str(args.get("text", ""))
            pyautogui.write(text, interval=0.01)
            return _text(f"typed {len(text)} characters.")
        except Exception as e:  # noqa: BLE001
            return _text(f"type_text failed: {e}")

    @tool("key", "Press a keyboard hotkey combo, e.g. 'ctrl+c' or 'win+d'. "
          "Combine keys with '+'.", {"keys": str})
    async def key(args: dict) -> dict:
        try:
            import pyautogui  # lazy
            raw = str(args.get("keys", "")).strip()
            if not raw:
                return _text("key failed: no keys given.")
            alias = {"win": "winleft", "cmd": "winleft", "windows": "winleft"}
            parts = [alias.get(p.strip().lower(), p.strip().lower())
                     for p in raw.split("+") if p.strip()]
            pyautogui.hotkey(*parts)
            return _text(f"pressed {'+'.join(parts)}.")
        except Exception as e:  # noqa: BLE001
            return _text(f"key failed: {e}")

    @tool("scroll", "Scroll the mouse wheel by the given amount (positive = "
          "up, negative = down).", {"amount": int})
    async def scroll(args: dict) -> dict:
        try:
            import pyautogui  # lazy
            amount = int(args["amount"])
            pyautogui.scroll(amount)
            return _text(f"scrolled {amount}.")
        except Exception as e:  # noqa: BLE001
            return _text(f"scroll failed: {e}")

    @tool("launch", "Open/launch/focus any app by name the way a person would "
          "from the Start menu (e.g. 'Discord', 'Marvel Rivals', 'Microsoft "
          "Teams', 'notepad'), or a path/URL.", {"app": str})
    async def launch(args: dict) -> dict:
        app = str(args.get("app", "")).strip()
        if not app:
            return _text("launch failed: no app given.")
        # Direct path / URL / .exe → open straight away.
        if (os.path.exists(app) or app.startswith(("http://", "https://"))
                or "\\" in app or "/" in app or app.lower().endswith(".exe")):
            try:
                os.startfile(app)  # type: ignore[attr-defined]
                return _text(f"launched {app}.")
            except Exception:  # noqa: BLE001
                pass
        # Otherwise resolve through the Start-menu index (Store/Steam/desktop).
        ok, matched = shell_launch(app)
        return _text(f"launched {matched}." if ok
                     else f"couldn't find an app matching '{app}' "
                          f"({matched}).")

    @tool("windows", "List visible top-level windows with their titles and "
          "bounds (left,top,right,bottom).", {})
    async def windows(args: dict) -> dict:
        try:
            import win32gui  # lazy (pywin32)
        except Exception:  # noqa: BLE001
            return _text("windows unavailable: pywin32 is not installed.")
        try:
            found: list[str] = []

            def _cb(hwnd, _):
                if not win32gui.IsWindowVisible(hwnd):
                    return
                title = win32gui.GetWindowText(hwnd)
                if not title.strip():
                    return
                l, t, r, b = win32gui.GetWindowRect(hwnd)
                found.append(f"{title} | {l},{t},{r},{b}")

            win32gui.EnumWindows(_cb, None)
            if not found:
                return _text("no visible titled windows found.")
            return _text("\n".join(found))
        except Exception as e:  # noqa: BLE001
            return _text(f"windows failed: {e}")

    @tool("snapshot", "See the REAL buttons: returns a screenshot AND a text "
          "list of the foreground window's UI elements (accessibility tree).",
          {})
    async def snapshot(args: dict) -> dict:
        content: list[dict] = []
        # 1) screenshot (downscaled, with the geometry/scale note)
        try:
            png, note = _grab_screen()
            data = base64.b64encode(png).decode("ascii")
            content.append({"type": "text", "text": note})
            content.append({"type": "image", "data": data,
                            "mimeType": "image/png"})
        except Exception as e:  # noqa: BLE001
            content.append({"type": "text",
                            "text": f"(screenshot failed: {e})"})
        # 2) accessibility tree of the foreground window
        try:
            import uiautomation as auto  # lazy
            fg = auto.GetForegroundControl()
            if fg is None:
                content.append({"type": "text",
                                "text": "No foreground window found."})
                return {"content": content}
            lines: list[str] = []

            def _walk(ctrl, depth):
                if depth > 2:
                    return
                try:
                    rect = ctrl.BoundingRectangle
                    bounds = (f"{rect.left},{rect.top},"
                              f"{rect.right},{rect.bottom}")
                except Exception:  # noqa: BLE001
                    bounds = "?"
                name = ""
                try:
                    name = ctrl.Name
                except Exception:  # noqa: BLE001
                    name = ""
                ctype = getattr(ctrl, "ControlTypeName", "Control")
                indent = "  " * depth
                lines.append(f"{indent}{ctype} | {name} | {bounds}")
                try:
                    for child in ctrl.GetChildren():
                        _walk(child, depth + 1)
                except Exception:  # noqa: BLE001
                    pass

            _walk(fg, 0)
            tree = "\n".join(lines) if lines else "(no elements read)"
            content.append({"type": "text",
                            "text": "Foreground UI elements "
                                    "(ControlType | Name | bounds):\n" + tree})
        except Exception:  # noqa: BLE001
            content.append({"type": "text",
                            "text": "Accessibility tree unavailable "
                                    "(uiautomation not installed)."})
        return {"content": content}

    @tool("shell", "Run a PowerShell command and return its output.",
          {"command": str})
    async def shell(args: dict) -> dict:
        try:
            command = str(args.get("command", ""))
            proc = subprocess.run(
                ["powershell", "-NoProfile", "-Command", command],
                capture_output=True, text=True, timeout=60,
            )
            out = (proc.stdout or "") + (proc.stderr or "")
            if len(out) > 4000:
                out = out[:4000] + "\n…(truncated)"
            return _text(out or "(no output)")
        except Exception as e:  # noqa: BLE001
            return _text(f"shell failed: {e}")

    return create_sdk_mcp_server(
        name="desktop",
        version="1.0.0",
        tools=[screenshot, click, move, type_text, key, scroll,
               launch, windows, snapshot, shell],
    )