#!/usr/bin/env python3
"""
Linux driver for octavus-computer-use MCP.

Implements the same CLI + JSON protocol as the macOS Swift driver,
using AT-SPI2 for accessibility tree, xdotool for input simulation,
scrot/xdpyinfo for screenshots, and Pillow for badge annotation.

Dependencies (apt):
    python3-gi gir1.2-atspi-2.0 at-spi2-core xdotool scrot python3-pil
    fonts-dejavu-core (a scalable bold TTF for grid/label overlays; see
    _BOLD_FONT_CANDIDATES)

Usage mirrors the macOS driver exactly:
    linux-driver.py observe --screenshot /tmp/shot.png [--pid PID]
    linux-driver.py click <x> <y>
    linux-driver.py double-click <x> <y>
    linux-driver.py right-click <x> <y>
    linux-driver.py type <text>
    linux-driver.py key <key-spec>
    linux-driver.py scroll <direction> [steps] [--at x y]
    linux-driver.py drag <from_x> <from_y> <to_x> <to_y> [--steps N] [--path x1,y1 ...]
    linux-driver.py screenshot <path>
    linux-driver.py screenshot-grid <path> --region x y w h --grid <2|3>
    linux-driver.py activate <app-name>
"""

import glob
import json
import os
import re
import subprocess
import sys
import time
import shutil

# ---------------------------------------------------------------------------
# Output helpers (match Swift driver JSON contract)
# ---------------------------------------------------------------------------

def print_success(path=None):
    resp = {"success": True}
    if path is not None:
        resp["path"] = path
    print(json.dumps(resp))

def exit_with_error(message):
    print(json.dumps({"error": message}))
    sys.exit(1)

def _run(cmd, check=True, capture=True):
    try:
        r = subprocess.run(
            cmd,
            capture_output=capture,
            text=True,
            timeout=15,
            check=check,
        )
        return r
    except FileNotFoundError:
        exit_with_error(f"Required command not found: {cmd[0]}")
    except subprocess.CalledProcessError as e:
        exit_with_error(f"Command failed: {' '.join(cmd)}\n{e.stderr or e.stdout or ''}")
    except subprocess.TimeoutExpired:
        exit_with_error(f"Command timed out: {' '.join(cmd)}")

# ---------------------------------------------------------------------------
# Screen geometry
# ---------------------------------------------------------------------------

def get_screen_size():
    """Return (width, height) of the root window via xdotool."""
    try:
        out = subprocess.check_output(
            ["xdotool", "getdisplaygeometry"], text=True, timeout=5
        ).strip()
        w, h = out.split()
        return int(w), int(h)
    except Exception:
        return 1920, 1080

# ---------------------------------------------------------------------------
# AT-SPI2 accessibility tree
# ---------------------------------------------------------------------------

CLICKABLE_ROLES = {
    "push button", "toggle button", "link", "menu item", "menu",
    "check box", "check menu item", "radio button", "radio menu item",
    "combo box", "text", "password text", "entry", "spin button",
    "slider", "tab", "page tab", "tool bar button",
}

CONTAINER_ROLES = {
    "frame", "dialog", "window", "scroll pane", "panel",
    "filler", "application", "desktop frame", "root pane",
    "layered pane", "split pane", "viewport",
}

ACTIONABLE_ACTIONS = {"click", "press", "activate", "toggle", "jump"}


def _try_import_atspi():
    """Import AT-SPI2 via GObject introspection and connect to the a11y bus."""
    try:
        import gi
        gi.require_version("Atspi", "2.0")
        from gi.repository import Atspi
        Atspi.init()
        return Atspi
    except Exception:
        return None


def _has_actionable(node, Atspi):
    try:
        action = node.get_action_iface()
        if action is None:
            return False
        n = action.get_n_actions()
        for i in range(n):
            name = action.get_action_name(i).lower()
            if name in ACTIONABLE_ACTIONS:
                return True
    except Exception:
        pass
    return False


def _overlaps_screen(x, y, w, h, screen_w, screen_h):
    """Whether a screen-coordinate box intersects the visible screen."""
    return x < screen_w and y < screen_h and x + w > 0 and y + h > 0


def _walk_atspi(
    node, elements, Atspi, screen_w, screen_h, depth=0, max_depth=20, trust_geometry=False
):
    if depth > max_depth:
        return

    visible = False
    showing = False
    try:
        state_set = node.get_state_set()
        visible = state_set.contains(Atspi.StateType.VISIBLE)
        showing = state_set.contains(Atspi.StateType.SHOWING)
    except Exception:
        pass

    rect = None
    try:
        comp = node.get_component_iface()
        if comp is not None:
            extents = comp.get_extents(Atspi.CoordType.SCREEN)
            rect = (extents.x, extents.y, extents.width, extents.height)
    except Exception:
        rect = None

    on_screen = (
        rect is not None
        and rect[2] > 0
        and rect[3] > 0
        and _overlaps_screen(rect[0], rect[1], rect[2], rect[3], screen_w, screen_h)
    )

    # Include this node as an actionable element when it is an on-screen clickable
    # widget. Inclusion keys on real screen geometry, not the AT-SPI SHOWING flag:
    # some toolkits (notably GTK4, e.g. gnome-text-editor) leave SHOWING unset on
    # every on-screen descendant below the top frame, so gating on SHOWING would
    # drop the entire app's controls. `trust_geometry` extends this to the VISIBLE
    # flag, which GTK4 likewise leaves unset on those descendants: for a native
    # foreground app, geometry + role decide inclusion; for Chromium (whose web
    # a11y flags are reliable) VISIBLE stays required, so web labeling never
    # surfaces hidden-but-positioned elements. Off-screen and zero-area nodes fail
    # the on_screen test regardless.
    if rect is not None and on_screen and (trust_geometry or visible):
        try:
            role_name = node.get_role_name().lower()
        except Exception:
            role_name = ""
        if role_name and (role_name in CLICKABLE_ROLES or _has_actionable(node, Atspi)):
            x, y, w, h = rect
            area = w * h
            screen_area = screen_w * screen_h
            too_big = (
                role_name in CONTAINER_ROLES and area > screen_area * 0.25
            ) or area > screen_area * 0.5
            if not too_big:
                title = ""
                try:
                    title = node.get_name() or ""
                except Exception:
                    pass
                if not title:
                    try:
                        title = node.get_description() or ""
                    except Exception:
                        pass
                if not title:
                    try:
                        text_iface = node.get_text_iface()
                        if text_iface:
                            val = text_iface.get_text(0, min(80, text_iface.get_character_count()))
                            if val:
                                title = val
                    except Exception:
                        pass
                if len(title) > 60:
                    title = title[:60] + "..."

                elements.append({
                    "role": role_name,
                    "title": title,
                    "state": _get_element_state(node, role_name, Atspi),
                    "x": x, "y": y,
                    "width": w, "height": h,
                })

    # Descend into children unless this node is provably off-screen. Keep walking
    # through nodes with no geometry (structural containers such as the
    # `application` node, which reports VISIBLE=False and no extents) and through
    # on-screen containers that merely lack SHOWING - prune only a node that has a
    # real box that is off-screen or zero-area and is also not SHOWING. That
    # bounds the walk over large web trees (off-viewport content) without severing
    # a native app's subtree at a SHOWING-less container.
    if rect is not None and not on_screen and not showing:
        return

    try:
        n_children = node.get_child_count()
        for i in range(n_children):
            child = node.get_child_at_index(i)
            if child is not None:
                _walk_atspi(
                    child, elements, Atspi, screen_w, screen_h, depth + 1, max_depth, trust_geometry
                )
    except Exception:
        pass


def _get_element_state(node, role_name, Atspi):
    parts = []
    try:
        state_set = node.get_state_set()
        if role_name in ("check box", "check menu item", "radio button", "radio menu item", "toggle button"):
            if state_set.contains(Atspi.StateType.CHECKED):
                parts.append("selected")
            else:
                parts.append("unselected")
        if state_set.contains(Atspi.StateType.FOCUSED):
            parts.append("focused")
        if role_name in ("text", "password text", "entry", "spin button", "combo box"):
            try:
                text_iface = node.get_text_iface()
                if text_iface:
                    val = text_iface.get_text(0, min(40, text_iface.get_character_count()))
                    if val:
                        if len(val) > 40:
                            val = val[:40] + "..."
                        parts.append(f'value="{val}"')
            except Exception:
                pass
    except Exception:
        pass
    return ", ".join(parts)


def _dedup_elements(elements, screen_w, screen_h):
    elements.sort(key=lambda e: e["width"] * e["height"])

    deduped = []
    for el in elements:
        area = el["width"] * el["height"]
        if area > screen_w * screen_h * 0.5:
            continue

        cx = el["x"] + el["width"] / 2
        cy = el["y"] + el["height"] / 2
        is_dup = False
        for ex in deduped:
            ex_cx = ex["x"] + ex["width"] / 2
            ex_cy = ex["y"] + ex["height"] / 2
            dist = ((cx - ex_cx) ** 2 + (cy - ex_cy) ** 2) ** 0.5
            if dist < 12.0:
                is_dup = True
                break
            if (ex["x"] >= el["x"] and ex["y"] >= el["y"]
                    and ex["x"] + ex["width"] <= el["x"] + el["width"]
                    and ex["y"] + ex["height"] <= el["y"] + el["height"]
                    and dist < 30.0):
                is_dup = True
                break
        if not is_dup:
            deduped.append(el)

    deduped.sort(key=lambda e: (round(e["y"] / 15) * 15, e["x"]))
    return deduped


def _get_active_window_pid():
    try:
        wid = subprocess.check_output(
            ["xdotool", "getactivewindow"], text=True, timeout=5
        ).strip()
        pid = subprocess.check_output(
            ["xdotool", "getwindowpid", wid], text=True, timeout=5
        ).strip()
        return int(pid)
    except Exception:
        return None


def _get_active_window_geometry():
    """Screen rect of the focused window as {x, y, width, height}, or None.

    Used to clip a recovered element set to the window the agent is actually
    looking at (one desktop, what-you-see-is-what-you-get).
    """
    try:
        out = subprocess.check_output(
            ["xdotool", "getactivewindow", "getwindowgeometry", "--shell"],
            text=True,
            timeout=5,
        )
    except Exception:
        return None

    vals = {}
    for line in out.splitlines():
        key, sep, value = line.partition("=")
        if sep:
            vals[key.strip()] = value.strip()
    try:
        return {
            "x": int(vals["X"]),
            "y": int(vals["Y"]),
            "width": int(vals["WIDTH"]),
            "height": int(vals["HEIGHT"]),
        }
    except (KeyError, ValueError):
        return None


def _active_window_is_chromium(pid):
    """Whether the foreground PID is a Chrome/Chromium process.

    Read from /proc rather than the window class: the cloud xdotool build has no
    getwindowclassname and xprop is not installed, but the browser owns the
    focused window's _NET_WM_PID, so its process name is the reliable signal.
    """
    if pid is None:
        return False
    for proc_file in (f"/proc/{pid}/comm", f"/proc/{pid}/cmdline"):
        try:
            with open(proc_file, "rb") as f:
                data = f.read().replace(b"\x00", b" ").decode("utf-8", "ignore").lower()
        except Exception:
            continue
        if "chrome" in data or "chromium" in data:
            return True
    return False


def _collect_elements(desktop, target_pid, Atspi, screen_w, screen_h, trust_geometry=False):
    """Collect actionable elements from the AT-SPI app(s) matching target_pid.

    Strictly PID-scoped: a None target_pid (foreground PID unavailable) matches no
    app and returns empty, so labeling never falls back to walking the whole
    desktop, which would overlay the always-open browser's labels onto whatever
    native app is in front. `trust_geometry` is forwarded to the walk so a native
    foreground app's on-screen controls are included from geometry even when its
    toolkit leaves the AT-SPI VISIBLE/SHOWING flags unset (GTK4).
    """
    elements = []
    n_apps = desktop.get_child_count()
    for i in range(n_apps):
        app = desktop.get_child_at_index(i)
        if app is None:
            continue
        try:
            if app.get_process_id() != target_pid:
                continue
        except Exception:
            continue
        _walk_atspi(app, elements, Atspi, screen_w, screen_h, trust_geometry=trust_geometry)
    return elements


def _collect_chromium(desktop, Atspi, screen_w, screen_h):
    """Walk only the Chrome/Chromium AT-SPI app node(s), matched by name.

    Recovery for when the foreground window is Chromium but the PID-scoped walk
    came back empty (a lazily-exported page tree, or an active-window PID that
    disagrees with the AT-SPI app PID). Scoping by name keeps the recovery on the
    browser alone, so a native app in front is never overlaid with the always-open
    browser's labels.
    """
    elements = []
    try:
        n_apps = desktop.get_child_count()
    except Exception:
        return elements
    for i in range(n_apps):
        try:
            app = desktop.get_child_at_index(i)
        except Exception:
            continue
        if app is None:
            continue
        try:
            name = (app.get_name() or "").lower()
        except Exception:
            name = ""
        if "chrome" in name or "chromium" in name:
            _walk_atspi(app, elements, Atspi, screen_w, screen_h)
    return elements


def _within_rect(elements, rect):
    """Keep only elements whose center lies inside the given screen rect."""
    if not rect:
        return elements
    rx, ry, rw, rh = rect["x"], rect["y"], rect["width"], rect["height"]
    kept = []
    for el in elements:
        cx = el["x"] + el["width"] / 2
        cy = el["y"] + el["height"] / 2
        if rx <= cx <= rx + rw and ry <= cy <= ry + rh:
            kept.append(el)
    return kept


def enumerate_elements(pid=None):
    Atspi = _try_import_atspi()
    if Atspi is None:
        exit_with_error(
            "AT-SPI2 not available. Install: sudo apt install python3-gi gir1.2-atspi-2.0 at-spi2-core"
        )

    screen_w, screen_h = get_screen_size()
    target_pid = pid or _get_active_window_pid()
    # A native foreground app (GTK4 etc.) leaves the AT-SPI VISIBLE/SHOWING flags
    # unset on its on-screen descendants, so trust geometry for it; Chromium's web
    # a11y flags are reliable, so keep requiring VISIBLE there.
    foreground_is_chromium = _active_window_is_chromium(target_pid)

    desktop = Atspi.get_desktop(0)
    elements = _collect_elements(
        desktop, target_pid, Atspi, screen_w, screen_h, trust_geometry=not foreground_is_chromium
    )

    # The foreground app exported its own actionable tree - use it verbatim. It is
    # already scoped to the single window the agent sees (one desktop, WYSIWYG).
    if not elements and foreground_is_chromium:
        # Chromium publishes its accessibility tree lazily, and its active-window
        # PID can disagree with the AT-SPI app PID, so the PID-scoped walk above
        # can miss it. Recover by walking only the Chrome app node(s), retrying
        # briefly for the lazy export, then clip to the focused window. We never
        # fall back to the whole desktop, which would overlay the always-open
        # browser's labels on top of whatever native app is in front.
        active_rect = _get_active_window_geometry()
        for _ in range(2):
            time.sleep(0.6)
            desktop = Atspi.get_desktop(0)
            elements = _collect_chromium(desktop, Atspi, screen_w, screen_h)
            if elements:
                break
        elements = _within_rect(elements, active_rect)

    # A native (non-Chromium) foreground app that exports nothing yields an empty
    # result on purpose: the agent falls back to screenshot_grid coordinates
    # rather than being handed the background browser's elements.

    deduped = _dedup_elements(elements, screen_w, screen_h)

    for idx, el in enumerate(deduped):
        el["number"] = idx + 1

    return deduped


# ---------------------------------------------------------------------------
# Screenshot capture
# ---------------------------------------------------------------------------

def _screenshot_tool_cmds(output_path):
    """Full-display capture commands tried in order. scrot is the primary (present
    on our images); the rest are fallbacks so a missing or broken scrot degrades
    gracefully instead of blinding the loop. Each writes ``output_path`` directly."""
    return (
        ["scrot", "-o", output_path],
        ["maim", output_path],
        ["import", "-window", "root", output_path],
        ["gnome-screenshot", "-f", output_path],
    )


def _pillow_x11_grab(output_path):
    """Last-resort capture via Pillow's native X11 grab (no external tool). Returns
    True on a non-empty PNG, else False."""
    try:
        from PIL import ImageGrab

        image = ImageGrab.grab(xdisplay=os.environ.get("DISPLAY") or ":0")
        image.save(output_path, "PNG")
        return os.path.exists(output_path) and os.path.getsize(output_path) > 0
    except Exception:
        return False


def capture_screenshot(output_path):
    """Take a full-display screenshot (X11).

    Prefers scrot; falls back to other common capture tools, then Pillow's native
    X11 grab, so a missing or failing scrot (e.g. it was never installed, or crashed
    with a libpng write error on a full disk) degrades gracefully instead of aborting
    the run with "Required command not found: scrot".
    """
    for cmd in _screenshot_tool_cmds(output_path):
        if shutil.which(cmd[0]) is None:
            continue
        try:
            subprocess.run(cmd, capture_output=True, text=True, timeout=20, check=True)
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired, OSError):
            continue
        if os.path.exists(output_path) and os.path.getsize(output_path) > 0:
            return
    if _pillow_x11_grab(output_path):
        return
    exit_with_error(
        "Screenshot failed: no working capture tool. Install one of scrot, maim, "
        "imagemagick (import), or gnome-screenshot."
    )


def crop_screenshot(screenshot_path, region):
    """Crop a screenshot to the given region dict {x, y, width, height}."""
    try:
        from PIL import Image
    except ImportError:
        exit_with_error(
            "Pillow not installed. Install: pip3 install Pillow (or sudo apt install python3-pil)"
        )
    img = Image.open(screenshot_path)
    img_w, img_h = img.size
    screen_w, screen_h = get_screen_size()
    scale_x = img_w / screen_w
    scale_y = img_h / screen_h

    left = int(region["x"] * scale_x)
    top = int(region["y"] * scale_y)
    right = int((region["x"] + region["width"]) * scale_x)
    bottom = int((region["y"] + region["height"]) * scale_y)

    left = max(0, min(left, img_w))
    top = max(0, min(top, img_h))
    right = max(left, min(right, img_w))
    bottom = max(top, min(bottom, img_h))

    cropped = img.crop((left, top, right, bottom))
    cropped.info = img.info
    cropped.save(screenshot_path, "PNG")


# Bold sans-serif TrueType candidates, in preference order. DejaVu is the
# baseline on both the E2B image and the cloud AMI; Liberation and FreeSans are
# fallbacks so a substrate missing DejaVu still renders readable, scalable
# overlays instead of PIL's fixed-size bitmap default (which ignores the
# requested size and makes grid/label text unreadable).
_BOLD_FONT_CANDIDATES = (
    "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
    "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
    "/usr/share/fonts/truetype/freefont/FreeSansBold.ttf",
)


def load_bold_font(size):
    """Load a scalable bold font at the requested size from a known path,
    falling back to PIL's default only when no TrueType font is available."""
    from PIL import ImageFont

    for path in _BOLD_FONT_CANDIDATES:
        try:
            return ImageFont.truetype(path, size)
        except Exception:
            continue
    return ImageFont.load_default()


def draw_grid_overlay(screenshot_path, grid_size):
    """Draw a numbered grid overlay on an existing screenshot.
    grid_size is 2 or 3 (for 2x2 or 3x3)."""
    try:
        from PIL import Image, ImageDraw
    except ImportError:
        exit_with_error(
            "Pillow not installed. Install: pip3 install Pillow (or sudo apt install python3-pil)"
        )

    img = Image.open(screenshot_path).convert("RGBA")
    overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(overlay, "RGBA")
    img_w, img_h = img.size
    cell_w = img_w / grid_size
    cell_h = img_h / grid_size
    cell_min = min(cell_w, cell_h)

    line_width = max(1, min(int(cell_min / 40), 4))
    for i in range(1, grid_size):
        x = int(i * cell_w)
        draw.line([(x, 0), (x, img_h)], fill=(230, 25, 25, 216), width=line_width)
    for i in range(1, grid_size):
        y = int(i * cell_h)
        draw.line([(0, y), (img_w, y)], fill=(230, 25, 25, 216), width=line_width)

    font_size = max(8, min(int(cell_min / 6), 60))
    font = load_bold_font(font_size)

    for row in range(grid_size):
        for col in range(grid_size):
            cx = int(col * cell_w + cell_w / 2)
            cy = int(row * cell_h + cell_h / 2)
            label = f"R{row + 1}C{col + 1}"
            bbox = draw.textbbox((0, 0), label, font=font)
            tw = bbox[2] - bbox[0]
            th = bbox[3] - bbox[1]
            draw.text(
                (cx - tw // 2, cy - th // 2),
                label,
                fill=(230, 25, 25, 255),
                font=font,
            )

    result = Image.alpha_composite(img, overlay).convert("RGB")
    result.info = img.info
    result.save(screenshot_path, "PNG")


# ---------------------------------------------------------------------------
# Screenshot annotation (red numbered badges, matching macOS style)
# ---------------------------------------------------------------------------

def annotate_screenshot(screenshot_path, elements):
    if not elements:
        return

    try:
        from PIL import Image, ImageDraw
    except ImportError:
        exit_with_error(
            "Pillow not installed. Install: pip3 install Pillow (or sudo apt install python3-pil)"
        )

    img = Image.open(screenshot_path).convert("RGBA")
    draw = ImageDraw.Draw(img, "RGBA")
    img_w, img_h = img.size
    screen_w, screen_h = get_screen_size()
    scale_x = img_w / screen_w
    scale_y = img_h / screen_h

    font = load_bold_font(12)
    font_small = load_bold_font(10)

    for el in elements:
        num_str = str(el["number"])
        badge_size = 30 if len(num_str) > 2 else 24
        active_font = font_small if len(num_str) > 2 else font

        lx = el["x"] * scale_x - 2
        ly = el["y"] * scale_y - 2

        if lx < 0:
            lx = 0
        if ly < 0:
            ly = 0
        if lx + badge_size > img_w:
            lx = img_w - badge_size
        if ly + badge_size > img_h:
            ly = img_h - badge_size

        cx = lx + badge_size / 2
        cy = ly + badge_size / 2
        r = badge_size / 2

        draw.ellipse(
            [cx - r, cy - r, cx + r, cy + r],
            fill=(230, 25, 25, 242),
            outline=(255, 255, 255, 255),
            width=2,
        )

        bbox = draw.textbbox((0, 0), num_str, font=active_font)
        tw = bbox[2] - bbox[0]
        th = bbox[3] - bbox[1]
        tx = cx - tw / 2
        ty = cy - th / 2
        draw.text((tx, ty), num_str, fill=(255, 255, 255, 255), font=active_font)

    img = img.convert("RGB")
    img.save(screenshot_path, "PNG")


# ---------------------------------------------------------------------------
# Input simulation via xdotool
# ---------------------------------------------------------------------------

# xdotool mousemove --sync blocks until the X server confirms the pointer
# reached the target. Under display load - rapid screenshots + grid overlays and
# heavy page rendering during a CAPTCHA-solving loop - that confirmation can
# stall for many seconds and blow the 15s command timeout, failing the click
# outright. Bound the sync wait; if it stalls (or errors), fall back to a
# non-sync move so the cursor still lands and the click proceeds - a
# slightly-less-synchronous move instead of a hard failure.
_MOUSEMOVE_SYNC_TIMEOUT = 3

def mouse_move(x, y):
    ix, iy = str(int(x)), str(int(y))
    try:
        subprocess.run(
            ["xdotool", "mousemove", "--sync", ix, iy],
            capture_output=True,
            text=True,
            timeout=_MOUSEMOVE_SYNC_TIMEOUT,
            check=True,
        )
        time.sleep(0.025)
        return
    except FileNotFoundError:
        exit_with_error("Required command not found: xdotool")
    except (subprocess.TimeoutExpired, subprocess.CalledProcessError):
        pass
    # Sync confirmation stalled - re-issue without --sync (returns immediately)
    # and give the async motion a slightly longer settle before the caller
    # clicks, since a non-sync move carries no server arrival confirmation.
    _run(["xdotool", "mousemove", ix, iy])
    time.sleep(0.1)

def mouse_click(x, y):
    mouse_move(x, y)
    _run(["xdotool", "click", "1"])

def mouse_double_click(x, y):
    mouse_move(x, y)
    _run(["xdotool", "click", "--repeat", "2", "--delay", "50", "1"])

def mouse_right_click(x, y):
    mouse_move(x, y)
    _run(["xdotool", "click", "3"])

def type_text(text):
    _run(["xdotool", "type", "--clearmodifiers", "--delay", "5", text])

# Mapping from our key-spec names to xdotool key names
_KEY_MAP = {
    "return": "Return", "enter": "Return", "tab": "Tab", "space": "space",
    "delete": "BackSpace", "backspace": "BackSpace", "forwarddelete": "Delete",
    "escape": "Escape", "esc": "Escape",
    "up": "Up", "down": "Down", "left": "Left", "right": "Right",
    "home": "Home", "end": "End", "pageup": "Prior", "pagedown": "Next",
    "f1": "F1", "f2": "F2", "f3": "F3", "f4": "F4", "f5": "F5", "f6": "F6",
    "f7": "F7", "f8": "F8", "f9": "F9", "f10": "F10", "f11": "F11", "f12": "F12",
}

_MOD_MAP = {
    "cmd": "super", "command": "super", "super": "super",
    "shift": "shift",
    "alt": "alt", "option": "alt", "opt": "alt",
    "ctrl": "ctrl", "control": "ctrl",
}

def press_key(key_spec):
    parts = key_spec.lower().split("+")
    mods = []
    key_name = ""

    for part in parts:
        if part in _MOD_MAP:
            mods.append(_MOD_MAP[part])
        else:
            key_name = part

    xdot_key = _KEY_MAP.get(key_name, key_name)

    combo = "+".join(mods + [xdot_key]) if mods else xdot_key
    _run(["xdotool", "key", "--clearmodifiers", combo])

def _interpolate_points(from_x, from_y, to_x, to_y, steps):
    """Generate linearly interpolated points between two coordinates."""
    points = []
    for i in range(1, steps + 1):
        t = i / steps
        x = from_x + (to_x - from_x) * t
        y = from_y + (to_y - from_y) * t
        points.append((x, y))
    return points


def mouse_drag(from_x, from_y, to_x, to_y, steps=10, waypoints=None):
    """Perform a smooth drag from one position to another.

    Uses xdotool mousedown/mousemove/mouseup with --sync to ensure
    Chrome recognizes the drag gesture. The 16ms sleep between moves
    matches 60fps frame timing.

    When waypoints are provided the last waypoint is the destination, so
    (to_x, to_y) is not appended to the segment list. The button is
    released in a finally clause so a failed mousemove can't leave the
    cursor stuck in the down state.
    """
    mouse_move(from_x, from_y)
    _run(["xdotool", "mousedown", "1"])
    try:
        if waypoints:
            segments = [(from_x, from_y)] + list(waypoints)
        else:
            segments = [(from_x, from_y), (to_x, to_y)]
        for seg_idx in range(len(segments) - 1):
            sx, sy = segments[seg_idx]
            ex, ey = segments[seg_idx + 1]
            for px, py in _interpolate_points(sx, sy, ex, ey, steps):
                _run(["xdotool", "mousemove", "--sync", str(int(px)), str(int(py))])
                time.sleep(0.016)
    finally:
        _run(["xdotool", "mouseup", "1"], check=False)


def do_scroll(direction, steps=3, at_x=None, at_y=None):
    if at_x is not None and at_y is not None:
        mouse_move(at_x, at_y)
        time.sleep(0.05)
    else:
        sw, sh = get_screen_size()
        mouse_move(sw // 2, sh // 2)
        time.sleep(0.05)

    button = "5" if direction == "down" else "4"
    for _ in range(steps):
        _run(["xdotool", "click", "--repeat", "3", button])
        time.sleep(0.016)


# ---------------------------------------------------------------------------
# Application activation
# ---------------------------------------------------------------------------

_DESKTOP_DIRS = [
    "/usr/share/applications",
    "/usr/local/share/applications",
    os.path.expanduser("~/.local/share/applications"),
]


def _resolve_from_desktop_files(name):
    """Search .desktop files for an app matching *name*.

    Checks the Name= and GenericName= fields (case-insensitive), plus the
    desktop file basename (e.g. "chromium" matches "chromium-browser.desktop").
    Returns (binaries, desktop_ids) - binary commands from Exec= lines and
    matching .desktop filenames for gtk-launch.
    """
    name_lower = name.lower()
    binaries = []
    desktop_ids = []

    for desktop_dir in _DESKTOP_DIRS:
        for filepath in glob.glob(os.path.join(desktop_dir, "*.desktop")):
            try:
                with open(filepath) as f:
                    content = f.read()
            except Exception:
                continue

            basename = os.path.splitext(os.path.basename(filepath))[0].lower()
            matched = name_lower in basename

            if not matched:
                for line in content.splitlines():
                    if line.startswith("Name=") or line.startswith("GenericName="):
                        value = line.split("=", 1)[1].strip().lower()
                        if value == name_lower:
                            matched = True
                            break

            if not matched:
                continue

            desktop_ids.append(os.path.basename(filepath))

            for line in content.splitlines():
                if line.startswith("Exec="):
                    exec_cmd = line.split("=", 1)[1].strip()
                    exec_cmd = re.sub(r"\s+%[fFuUdDnNickvm]", "", exec_cmd)
                    binary = exec_cmd.split()[0]
                    if binary not in binaries:
                        binaries.append(binary)
                    break

    return binaries, desktop_ids


# WM_CLASS regex fragments for apps whose window title does not match the app name
# the agent asks to activate - notably Chrome `--app=<url>` windows, whose title is
# the page's, not "chrome". "chrome" matches the "google-chrome" class; "chromium"
# is listed explicitly since it does not contain "chrome".
_WINDOW_CLASS_ALIASES = {
    "chrome": ["chrome", "chromium"],
    "google chrome": ["chrome", "chromium"],
    "google-chrome": ["chrome", "chromium"],
    "chromium": ["chromium", "chrome"],
    "browser": ["chrome", "chromium", "firefox"],
    "firefox": ["firefox"],
}


def activate_app(name):
    def find_windows(flag, query):
        try:
            r = subprocess.run(
                ["xdotool", "search", flag, query],
                capture_output=True, text=True, timeout=5,
            )
            return [wid for wid in r.stdout.strip().split("\n") if wid]
        except Exception:
            return []

    def activate_existing_window():
        # Match an already-open window by title (--name) and by window class
        # (--class). A Chrome `--app=<url>` window carries the page title, not the
        # app name, so a title-only search misses it; the class aliases below focus
        # the running browser window instead of launching a second, duplicate one.
        attempts = [("--name", name), ("--name", name.lower())]
        attempts += [("--class", cls) for cls in _WINDOW_CLASS_ALIASES.get(name.lower(), [])]
        for flag, query in attempts:
            wids = find_windows(flag, query)
            if wids:
                # No --sync: `windowactivate --sync` blocks until the WM confirms
                # activation, which never arrives for a borderless/maximized
                # window under fluxbox/xfwm and hangs until the subprocess
                # timeout. A plain activate plus a short settle is enough; treat a
                # non-zero exit as best-effort so it can't abort labeling.
                _run(["xdotool", "windowactivate", wids[0]], check=False)
                time.sleep(0.2)
                return True
        return False

    def wait_for_window(timeout_s=4.0):
        deadline = time.time() + timeout_s
        while time.time() < deadline:
            if activate_existing_window():
                return True
            time.sleep(0.2)
        return False

    if activate_existing_window():
        return

    desktop_binaries, desktop_ids = _resolve_from_desktop_files(name)

    candidates = list(desktop_binaries)
    direct = name.lower()
    if direct not in candidates:
        candidates.append(direct)
    if name != direct and name not in candidates:
        candidates.append(name)

    launch_attempted = False
    for command in candidates:
        executable = shutil.which(command)
        if executable is None:
            continue

        launch_attempted = True
        try:
            subprocess.Popen(
                [executable],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            if wait_for_window():
                return
        except Exception:
            pass

    if not launch_attempted and desktop_ids:
        gtk_launch = shutil.which("gtk-launch")
        if gtk_launch:
            for desktop_id in desktop_ids:
                try:
                    subprocess.Popen(
                        [gtk_launch, desktop_id],
                        stdout=subprocess.DEVNULL,
                        stderr=subprocess.DEVNULL,
                    )
                    if wait_for_window():
                        return
                except Exception:
                    pass

    try:
        active_pid = _get_active_window_pid()
    except Exception:
        active_pid = None

    raise RuntimeError(
        f"Could not open or focus application '{name}'. "
        f"Active window pid remains {active_pid if active_pid is not None else 'unknown'}."
    )


# ---------------------------------------------------------------------------
# Main CLI
# ---------------------------------------------------------------------------

def main():
    args = sys.argv[1:]
    if not args:
        exit_with_error("Usage: linux-driver.py <command> [args...]")

    command = args[0]

    if command == "observe":
        pid = None
        screenshot_path = "/tmp/octavus-screenshot.png"
        i = 1
        while i < len(args):
            if args[i] == "--pid" and i + 1 < len(args):
                pid = int(args[i + 1])
                i += 2
            elif args[i] == "--screenshot" and i + 1 < len(args):
                screenshot_path = args[i + 1]
                i += 2
            else:
                i += 1

        elements = enumerate_elements(pid=pid)
        capture_screenshot(screenshot_path)
        annotate_screenshot(screenshot_path, elements)

        result = {
            "elements": elements,
            "screenshotPath": screenshot_path,
        }
        print(json.dumps(result, indent=2))

    elif command == "click":
        if len(args) < 3:
            exit_with_error("Usage: linux-driver.py click <x> <y>")
        mouse_click(float(args[1]), float(args[2]))
        print_success()

    elif command == "double-click":
        if len(args) < 3:
            exit_with_error("Usage: linux-driver.py double-click <x> <y>")
        mouse_double_click(float(args[1]), float(args[2]))
        print_success()

    elif command == "right-click":
        if len(args) < 3:
            exit_with_error("Usage: linux-driver.py right-click <x> <y>")
        mouse_right_click(float(args[1]), float(args[2]))
        print_success()

    elif command == "type":
        if len(args) < 2:
            exit_with_error("Usage: linux-driver.py type <text>")
        text = " ".join(args[1:])
        type_text(text)
        print_success()

    elif command == "key":
        if len(args) < 2:
            exit_with_error("Usage: linux-driver.py key <key-spec>")
        press_key(args[1])
        print_success()

    elif command == "scroll":
        if len(args) < 2:
            exit_with_error("Usage: linux-driver.py scroll <direction> [steps] [--at x y]")
        direction = args[1]
        steps = 3
        at_x = None
        at_y = None
        si = 2
        while si < len(args):
            if args[si] == "--at" and si + 2 < len(args):
                at_x = float(args[si + 1])
                at_y = float(args[si + 2])
                si += 3
            else:
                try:
                    steps = int(args[si])
                except ValueError:
                    pass
                si += 1
        do_scroll(direction, steps, at_x, at_y)
        print_success()

    elif command == "drag":
        if len(args) < 5:
            exit_with_error(
                "Usage: linux-driver.py drag <from_x> <from_y> <to_x> <to_y> "
                "[--steps N] [--path x1,y1 ...]"
            )
        from_x, from_y = float(args[1]), float(args[2])
        to_x, to_y = float(args[3]), float(args[4])
        steps = 10
        waypoints = None
        di = 5
        while di < len(args):
            if args[di] == "--steps" and di + 1 < len(args):
                steps = int(args[di + 1])
                di += 2
            elif args[di] == "--path":
                waypoints = []
                di += 1
                while di < len(args) and not args[di].startswith("--"):
                    parts = args[di].split(",")
                    waypoints.append((float(parts[0]), float(parts[1])))
                    di += 1
            else:
                di += 1
        mouse_drag(from_x, from_y, to_x, to_y, steps, waypoints)
        print_success()

    elif command == "screenshot":
        if len(args) < 2:
            exit_with_error("Usage: linux-driver.py screenshot <path>")
        path = args[1]
        capture_screenshot(path)
        print_success(path=path)

    elif command == "screenshot-grid":
        if len(args) < 2:
            exit_with_error(
                "Usage: linux-driver.py screenshot-grid <path> --region x y w h --grid <2|3>"
            )
        path = args[1]
        region = None
        grid_size = 3
        si = 2
        while si < len(args):
            if args[si] == "--region" and si + 4 < len(args):
                region = {
                    "x": float(args[si + 1]), "y": float(args[si + 2]),
                    "width": float(args[si + 3]), "height": float(args[si + 4]),
                }
                si += 5
            elif args[si] == "--grid" and si + 1 < len(args):
                grid_size = int(args[si + 1])
                si += 2
            else:
                si += 1
        if not region:
            exit_with_error("screenshot-grid requires --region x y w h")
        capture_screenshot(path)
        crop_screenshot(path, region)
        draw_grid_overlay(path, grid_size)
        print_success(path=path)

    elif command == "activate":
        if len(args) < 2:
            exit_with_error("Usage: linux-driver.py activate <app-name>")
        app_name = " ".join(args[1:])
        try:
            activate_app(app_name)
        except RuntimeError as e:
            exit_with_error(str(e))
        time.sleep(0.5)
        print_success()

    else:
        exit_with_error(f"Unknown command: '{command}'")


if __name__ == "__main__":
    main()