#!/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 _walk_atspi(node, elements, Atspi, screen_w, screen_h, depth=0, max_depth=20):
    if depth > max_depth:
        return
    try:
        state_set = node.get_state_set()
        if not state_set.contains(Atspi.StateType.VISIBLE):
            return
        if not state_set.contains(Atspi.StateType.SHOWING):
            return
    except Exception:
        pass

    try:
        comp = node.get_component_iface()
        if comp is not None:
            rect = comp.get_extents(Atspi.CoordType.SCREEN)
            x, y, w, h = rect.x, rect.y, rect.width, rect.height
            if w > 0 and h > 0:
                role_name = node.get_role_name().lower()
                is_click = role_name in CLICKABLE_ROLES or _has_actionable(node, Atspi)

                if is_click:
                    area = w * h
                    screen_area = screen_w * screen_h
                    if role_name in CONTAINER_ROLES and area > screen_area * 0.25:
                        pass
                    elif area > screen_area * 0.5:
                        pass
                    else:
                        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] + "..."

                        state = _get_element_state(node, role_name, Atspi)

                        elements.append({
                            "role": role_name,
                            "title": title,
                            "state": state,
                            "x": x, "y": y,
                            "width": w, "height": h,
                        })
    except Exception:
        pass

    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)
    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 _collect_elements(desktop, target_pid, Atspi, screen_w, screen_h):
    elements = []
    matched_target = False

    n_apps = desktop.get_child_count()
    for i in range(n_apps):
        app = desktop.get_child_at_index(i)
        if app is None:
            continue

        if target_pid is not None:
            try:
                app_pid = app.get_process_id()
                if app_pid != target_pid:
                    continue
                matched_target = True
            except Exception:
                continue

        _walk_atspi(app, elements, Atspi, screen_w, screen_h)

    return elements, matched_target


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()

    desktop = Atspi.get_desktop(0)
    elements, matched_target = _collect_elements(
        desktop, target_pid, Atspi, screen_w, screen_h
    )

    # Linux/X11 accessibility can disagree with xdotool's active-window PID,
    # especially for Chromium/Electron renderers. If the focused PID yields
    # nothing, retry against the full AT-SPI desktop tree.
    if target_pid is not None and (not matched_target or not elements):
        elements, _ = _collect_elements(desktop, None, Atspi, screen_w, screen_h)

    # Chromium publishes its accessibility tree to AT-SPI lazily, so a label fired
    # right after navigation can beat that export and see an empty desktop. The
    # a11y bus is already enabled by the runtime at startup, so wait briefly and
    # re-walk the full tree once before giving up.
    if not elements:
        time.sleep(0.6)
        desktop = Atspi.get_desktop(0)
        elements, _ = _collect_elements(desktop, None, Atspi, screen_w, screen_h)

    deduped = _dedup_elements(elements, screen_w, screen_h)

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

    return deduped


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

def capture_screenshot(output_path):
    """Take a screenshot using scrot (X11)."""
    _run(["scrot", "-o", output_path])


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


def activate_app(name):
    def find_windows(query):
        try:
            r = subprocess.run(
                ["xdotool", "search", "--name", 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():
        for query in (name, name.lower()):
            wids = find_windows(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()