vscode_agent/vscode_controller.py

"""
vscode_controller.py
====================
Handles VS Code-specific automation:
  - Create / open files and folders
  - Write, append, or prepend code into files
  - Run commands in the integrated terminal
  - Install / list extensions

Dependencies: code CLI, wmctrl, xdotool
"""

import os
import time
import subprocess
from pathlib import Path


# ── File operations ───────────────────────────────────────────────────────────

def open_in_vscode(path: str, new_window: bool = False) -> dict:
    """Open a file or folder in VS Code."""
    expanded = os.path.expanduser(path)
    if not os.path.exists(expanded):
        return {"success": False, "message": f"Path not found: '{path}'", "action": "error"}

    cmd = ["code"] + (["--new-window"] if new_window else []) + [expanded]
    subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                     start_new_session=True)
    return {"success": True, "message": f"Opened '{path}' in VS Code.",
            "action": "opened", "path": expanded}


def create_file(filepath: str, content: str = "", open_after: bool = True) -> dict:
    # ── Ensure absolute path ──────────────────────────────────────────────
    if not filepath or not filepath.strip():
        return {"success": False, "message": "No file path provided.", "action": "error"}

    expanded = os.path.expanduser(filepath)

    # If still not absolute, reject — don't guess
    if not os.path.isabs(expanded):
        return {"success": False,
                "message": f"Could not resolve path: '{filepath}'. Please provide a full path.",
                "action": "error"}

    try:
        Path(expanded).parent.mkdir(parents=True, exist_ok=True)
        Path(expanded).write_text(content, encoding="utf-8")

        # ── Verify file actually exists on disk before opening ────────────
        if not Path(expanded).exists():
            return {"success": False, "message": f"File write failed: '{expanded}'", "action": "error"}

        if open_after:
            subprocess.Popen(
                ["code", "--goto", expanded],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                start_new_session=True,
            )

        return {"success": True, "message": f"Created '{expanded}'.",
                "action": "created", "path": expanded}

    except Exception as e:
        return {"success": False, "message": f"Failed to create file: {e}", "action": "error"}


def create_folder(folderpath: str, open_after: bool = True) -> dict:
    """Create a folder and open it as a VS Code workspace."""
    expanded = os.path.expanduser(folderpath)
    try:
        Path(expanded).mkdir(parents=True, exist_ok=True)
        if open_after:
            subprocess.Popen(["code", expanded],
                             stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                             start_new_session=True)
        return {"success": True, "message": f"Created folder '{folderpath}'.",
                "action": "folder_created", "path": expanded}
    except Exception as e:
        return {"success": False, "message": f"Failed to create folder: {e}", "action": "error"}


def write_to_file(filepath: str, content: str, mode: str = "overwrite") -> dict:
    """
    Write content to an existing file.
    mode: 'overwrite' | 'append' | 'prepend'
    """
    expanded = os.path.expanduser(filepath)
    p = Path(expanded)
    if not p.exists():
        return {"success": False, "message": f"File not found: '{filepath}'", "action": "error"}

    try:
        if mode == "overwrite":
            p.write_text(content, encoding="utf-8")
        elif mode == "append":
            with open(expanded, "a", encoding="utf-8") as f:
                f.write("\n" + content)
        elif mode == "prepend":
            existing = p.read_text(encoding="utf-8")
            p.write_text(content + "\n" + existing, encoding="utf-8")
        else:
            return {"success": False, "message": f"Unknown write mode: '{mode}'", "action": "error"}

        return {"success": True, "message": f"Written to '{filepath}' ({mode}).",
                "action": "written", "path": expanded, "mode": mode}
    except Exception as e:
        return {"success": False, "message": f"Failed to write: {e}", "action": "error"}


# ── Terminal operations ───────────────────────────────────────────────────────

def run_in_terminal(command: str) -> dict:
    """
    Focus VS Code, open the integrated terminal, and run a command.
    Uses wmctrl + xdotool.
    """
    try:
        # 1. Focus VS Code window
        subprocess.run(["wmctrl", "-x", "-a", "code.Code"],
                       capture_output=True, timeout=5)
        time.sleep(0.6)

        # 2. Open integrated terminal (Ctrl + `)
        subprocess.run(["xdotool", "key", "ctrl+grave"], timeout=5)
        time.sleep(0.8)

        # 3. Type command + Enter
        subprocess.run(
            ["xdotool", "type", "--clearmodifiers", "--delay", "40", command],
            timeout=15
        )
        time.sleep(0.2)
        subprocess.run(["xdotool", "key", "Return"], timeout=5)

        return {"success": True, "message": f"Ran in terminal: `{command}`",
                "action": "terminal_run", "command": command}
    except Exception as e:
        return {"success": False, "message": f"Terminal run failed: {e}", "action": "error"}


# ── Extension operations ──────────────────────────────────────────────────────

def install_extension(extension_id: str) -> dict:
    """Install a VS Code extension by its marketplace ID (e.g. 'ms-python.python')."""
    try:
        result = subprocess.run(
            ["code", "--install-extension", extension_id, "--force"],
            capture_output=True, text=True, timeout=120
        )
        if result.returncode == 0:
            return {"success": True,
                    "message": f"Installed extension '{extension_id}'.",
                    "action": "extension_installed", "extension": extension_id}
        return {"success": False, "message": result.stderr.strip() or "Install failed.",
                "action": "error"}
    except Exception as e:
        return {"success": False, "message": f"Extension install failed: {e}", "action": "error"}


def list_extensions() -> dict:
    """List all installed VS Code extensions."""
    try:
        result = subprocess.run(["code", "--list-extensions"],
                                capture_output=True, text=True, timeout=15)
        exts = [e.strip() for e in result.stdout.strip().splitlines() if e.strip()]
        return {"success": True, "extensions": exts, "count": len(exts)}
    except Exception as e:
        return {"success": False, "message": str(e), "extensions": []}
    

def get_open_vscode_windows() -> list[dict]:
    """
    Returns list of open VS Code windows with their workspace names.
    Each: {"id": "0x04000004", "title": "AI OS", "full_title": "main_agent.py - AI OS - Visual Studio Code"}
    """
    try:
        result = subprocess.run(
            ["wmctrl", "-l", "-x"],
            capture_output=True, text=True, timeout=5
        )
        windows = []
        for line in result.stdout.strip().splitlines():
            parts = line.split(None, 4)
            if len(parts) < 5:
                continue
            win_class = parts[2].lower()
            title     = parts[4].strip()
            if "code.code" in win_class or "visual studio code" in title.lower():
                # Extract workspace name — "main_agent.py - AI OS - Visual Studio Code"
                # → "AI OS"
                segments = [s.strip() for s in title.split(" - ")]
                workspace = segments[-2] if len(segments) >= 2 else title
                windows.append({
                    "id":         parts[0],
                    "title":      workspace,
                    "full_title": title,
                })
        return windows
    except Exception as e:
        print(f"[VSCodeController] Window list failed: {e}")
        return []


def focus_vscode_window(window_id: str) -> dict:
    """Focus a specific VS Code window by its wmctrl ID."""
    try:
        subprocess.run(
            ["wmctrl", "-i", "-a", window_id],
            capture_output=True, timeout=5
        )
        return {"success": True, "message": f"Focused window {window_id}"}
    except Exception as e:
        return {"success": False, "message": str(e)}