vision_agent.py

"""
Vision-First Web Automation Agent
Uses GPT-4V to analyze screenshots and guide navigation step-by-step.
"""

import json
import base64
from typing import Dict, List, Any, Optional
from openai import OpenAI
from config import Config
from utils import setup_logger

logger = setup_logger(__name__, Config.LOG_FILE, Config.LOG_LEVEL)


class VisionTaskInterpreter:
    """
    Vision-first task interpreter that uses GPT-4V to:
    1. Analyze current UI state from screenshots
    2. Determine next action to take
    3. Guide navigation step-by-step
    """

    def __init__(self):
        """Initialize with OpenAI client for GPT-4V."""
        if not Config.OPENAI_API_KEY:
            raise ValueError("OPENAI_API_KEY is required for vision-guided automation")

        self.client = OpenAI(api_key=Config.OPENAI_API_KEY)
        self.model = "gpt-4o"  # Updated model with vision capabilities

    def parse_task(self, query: str) -> Dict[str, Any]:
        """
        Parse user query to extract app name and task goal.

        Args:
            query: Natural language query

        Returns:
            Task dictionary with app, goal, and description
        """
        prompt = f"""Parse this user query and extract the application name and task goal.

Query: "{query}"

Return ONLY a JSON object:
{{
  "app": "application name (Linear, Notion, etc.)",
  "task": "short task name (e.g., create project, add page)",
  "goal": "detailed goal description",
  "description": "{query}"
}}"""

        try:
            response = self.client.chat.completions.create(
                model="gpt-4o-mini",  # Use mini for simple text parsing
                messages=[{"role": "user", "content": prompt}],
                max_tokens=200
            )

            result_text = response.choices[0].message.content.strip()

            # Clean JSON if wrapped in markdown
            if "```json" in result_text:
                result_text = result_text.split("```json")[1].split("```")[0]
            elif "```" in result_text:
                result_text = result_text.split("```")[1].split("```")[0]

            task = json.loads(result_text.strip())
            logger.info(f"Parsed task: {task}")
            return task

        except Exception as e:
            logger.error(f"Error parsing task: {e}")
            # Fallback parsing
            return {
                "app": "Linear" if "linear" in query.lower() else "Unknown",
                "task": "navigate",
                "goal": query,
                "description": query
            }

    def get_next_action_with_som(
        self,
        labeled_screenshot_path: str,
        elements: List[Dict[str, Any]],
        task_goal: str,
        current_url: str,
        previous_actions: List[Dict[str, Any]]
    ) -> Dict[str, Any]:
        """
        Analyze labeled screenshot (with Set-of-Marks) and determine next action.
        This is the preferred method as it provides exact element IDs.

        Args:
            labeled_screenshot_path: Path to screenshot with numbered labels
            elements: List of interactive elements with IDs and positions
            task_goal: What we're trying to accomplish
            current_url: Current page URL
            previous_actions: List of actions taken so far

        Returns:
            Next action dictionary with element_id, action, value, reasoning
        """
        try:
            # Read and encode labeled screenshot
            with open(labeled_screenshot_path, 'rb') as f:
                screenshot_bytes = f.read()
            screenshot_base64 = base64.b64encode(screenshot_bytes).decode('utf-8')

            # Format elements for prompt
            elements_text = self._format_elements_for_prompt(elements)

            # Build context from previous actions
            actions_context = ""
            if previous_actions:
                actions_context = "Previous actions taken:\n"
                for i, action in enumerate(previous_actions[-5:], 1):  # Last 5 actions
                    action_desc = f"{action.get('action')}"
                    if action.get('element_id'):
                        action_desc += f" element #{action.get('element_id')}"
                    if action.get('description'):
                        action_desc += f" - {action.get('description')}"
                    actions_context += f"{i}. {action_desc}\n"

            url_context = f"Current URL: {current_url}\n" if current_url else ""

            prompt = f"""You are analyzing a web application screenshot with NUMBERED RED CIRCLES on interactive elements.

GOAL: {task_goal}

{url_context}{actions_context}

INTERACTIVE ELEMENTS (numbered on screenshot):
{elements_text}

INSTRUCTIONS:
- The screenshot shows red numbered circles (1, 2, 3, ...) on clickable elements
- Use BOTH the visual appearance AND the element descriptions above
- You MUST return the element_id number for click/fill/press actions
- Choose the element that best helps accomplish the goal

Available actions:
- click: Click an element (MUST provide element_id)
- fill: Type text into an element (MUST provide element_id and value)
- press: Press a keyboard key (MUST provide key)
- navigate: Go to a URL (MUST provide full url)
- wait: Wait for page to load (duration in ms)
- done: Task is complete

Return ONLY a JSON object:
{{
  "element_id": <number or null>,
  "action": "click|fill|press|navigate|wait|done",
  "value": "<text to type if filling>",
  "url": "<full URL if navigating>",
  "key": "<key name if pressing>",
  "reasoning": "why this helps accomplish the goal",
  "confidence": 0.0-1.0
}}

CRITICAL RULES:
- For click/fill actions, element_id is REQUIRED (must be a number from the list above)
- For navigate actions, url is REQUIRED (must be full URL like https://linear.app/team/projects)
- If you can't find the right element, return action="wait" with duration=2000
- If goal is accomplished, return action="done"

FORM COMPLETION RULES (VERY IMPORTANT):
- When you see a form or modal with input fields (name, title, description, etc.), you MUST fill the required fields BEFORE clicking submit/create/save buttons
- Look for empty input fields, text areas, or fields with placeholder text - these need to be filled first
- Common patterns: "Create", "Submit", "Save", "Add" buttons should only be clicked AFTER filling all necessary fields
- If a modal just opened, check for input fields before clicking any action buttons
- Fill fields with meaningful, descriptive values (e.g., "My New Project", "Test Task", etc.)

NOTION SLASH COMMANDS (CRITICAL FOR NOTION):
- Notion uses slash commands to insert blocks
- **CRITICAL**: Slash commands ONLY work in body content blocks, NOT in the title!
- **USE THE CORRECT SLASH COMMAND FOR THE TASK**:
  - For creating a DATABASE → use "/database" command
  - For creating a TABLE (simple table) → use "/table" command
  - For creating a HEADING → use "/heading" command
  - For creating a LIST → use "/bullet" or "/numbered" command
  - Match the slash command to the task goal!

**IMPORTANT NOTION WORKFLOW - CREATING BODY CONTENT BLOCKS**:
When you first create a new Notion page:
1. The page starts with ONLY a title field (H1 element) - there is NO body content block yet
2. **YOU MUST PRESS ENTER** from the title to create the first body content block below
3. Only AFTER pressing Enter will a body content block appear where you can type slash commands
4. Correct workflow: Click title → Press "Enter" key → Type "/database" → Select from menu

**HOW TO IDENTIFY TITLE vs BODY**:
- **NOTION_ROLE Metadata**: Some elements have a **NOTION_ROLE** field:
  - **NOTION_ROLE: TITLE** = Page title field (usually H1 tag) - **NEVER** type slash commands here!
  - **NOTION_ROLE: BODY** = Body content field - **ALWAYS** use this for slash commands!
- **Tag name**: H1 elements are always the title in Notion
- **Position**: Title is at the TOP (smaller y value ~100-200), body is BELOW (larger y value > 200)
- **DataPlaceholder**: Body areas may have "Type '/' for commands" text

**CRITICAL RULES**:
- **NEVER** type slash commands in an H1 element or element with **NOTION_ROLE: TITLE**
- **ALWAYS** press Enter after clicking the title to create a body content block
- If you only see title elements and no body elements, you MUST press "Enter" first
- **MATCH the slash command to the GOAL**:
  - If goal mentions "database" → type "/database"
  - If goal mentions "table" (and NOT database) → type "/table"
  - If goal mentions "heading" → type "/heading"
  - Look at the GOAL above to determine the correct slash command!
- Correct workflow: 1) Click title (or skip), 2) Press "Enter" key, 3) Type appropriate slash command based on GOAL, 4) Select from menu
"""

            # Call GPT-4o
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": prompt
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/png;base64,{screenshot_base64}",
                                    "detail": "high"
                                }
                            }
                        ]
                    }
                ],
                max_tokens=500
            )

            result_text = response.choices[0].message.content.strip()

            # Clean JSON if wrapped in markdown
            if "```json" in result_text:
                result_text = result_text.split("```json")[1].split("```")[0]
            elif "```" in result_text:
                result_text = result_text.split("```")[1].split("```")[0]

            next_action = json.loads(result_text.strip())

            # Validate element_id if provided
            if next_action.get('element_id'):
                element_id = next_action['element_id']
                if not any(el['id'] == element_id for el in elements):
                    logger.warning(f"GPT-4o returned invalid element_id: {element_id}")
                    next_action['element_id'] = None

            logger.info(f"GPT-4o SoM action: {next_action['action']} - element #{next_action.get('element_id', 'N/A')}")
            logger.info(f"Reasoning: {next_action['reasoning']}")
            logger.info(f"Confidence: {next_action.get('confidence', 'N/A')}")

            return next_action

        except Exception as e:
            logger.error(f"Error getting next action with SoM: {e}")
            # Return wait action as fallback
            return {
                "element_id": None,
                "action": "wait",
                "duration": 1000,
                "reasoning": f"Error analyzing screenshot: {str(e)}",
                "confidence": 0.0
            }

    def _format_elements_for_prompt(self, elements: List[Dict[str, Any]]) -> str:
        """Format elements list for GPT-4o prompt."""
        lines = []
        for el in elements:
            parts = [f"[{el['id']}]"]

            if el.get('text'):
                parts.append(f"Text: '{el['text'][:50]}'")

            if el.get('ariaLabel'):
                parts.append(f"Label: '{el['ariaLabel']}'")

            if el.get('placeholder'):
                parts.append(f"Placeholder: '{el['placeholder']}'")

            if el.get('dataPlaceholder'):
                parts.append(f"DataPlaceholder: '{el['dataPlaceholder']}'")

            if el.get('role'):
                parts.append(f"Role: {el['role']}")

            if el.get('tag'):
                parts.append(f"Tag: {el['tag']}")

            if el.get('type'):
                parts.append(f"Type: {el['type']}")

            if el.get('classes'):
                # Only show first few class names to keep prompt concise
                classes_preview = ' '.join(el['classes'].split()[:3])
                if classes_preview:
                    parts.append(f"Classes: {classes_preview}")

            if el.get('notion_role'):
                parts.append(f"**NOTION_ROLE: {el['notion_role']}**")

            lines.append(", ".join(parts))

        return "\n".join(lines)

    def get_next_action(
        self,
        screenshot_path: str,
        task_goal: str,
        previous_actions: List[Dict[str, Any]],
        current_url: str = ""
    ) -> Dict[str, Any]:
        """
        Analyze screenshot and determine the next action to take.

        Args:
            screenshot_path: Path to current screenshot
            task_goal: What we're trying to accomplish
            previous_actions: List of actions taken so far
            current_url: Current page URL

        Returns:
            Next action dictionary with action, target, value, reasoning
        """
        try:
            # Read and encode screenshot
            with open(screenshot_path, 'rb') as f:
                screenshot_bytes = f.read()
            screenshot_base64 = base64.b64encode(screenshot_bytes).decode('utf-8')

            # Build context from previous actions
            actions_context = ""
            if previous_actions:
                actions_context = "Previous actions taken:\n"
                for i, action in enumerate(previous_actions[-3:], 1):  # Last 3 actions
                    actions_context += f"{i}. {action.get('action')} - {action.get('description', 'N/A')}\n"

            url_context = f"Current URL: {current_url}\n" if current_url else ""

            prompt = f"""You are analyzing a web application screenshot to guide automated navigation.

GOAL: {task_goal}

{url_context}{actions_context}

Analyze the screenshot and determine the NEXT SINGLE ACTION to take.

Available actions:
- navigate: Go to a specific URL
- click: Click a button, link, or element
- fill: Type text into an input field or contenteditable element
- press: Press a keyboard key
- wait: Wait for page to load
- done: Task is complete

Return ONLY a JSON object:
{{
  "action": "action type",
  "target": "specific element description (REQUIRED for click/fill actions)",
  "value": "text to type (REQUIRED for fill actions)",
  "key": "key to press (REQUIRED for press actions)",
  "url": "FULL URL to navigate to (REQUIRED for navigate actions - MUST be complete URL like https://linear.app/team/issues)",
  "reasoning": "why this action will help accomplish the goal",
  "expected_result": "what should happen after this action",
  "confidence": 0.0-1.0
}}

CRITICAL: When action="navigate", you MUST provide the complete "url" field. Look at the current URL pattern to infer the correct URL.

IMPORTANT:
- Be VERY specific about element descriptions (use visible text, position, styling cues)
- If you see a form, identify ALL fields that need to be filled
- For create/add tasks, look for "New", "Add", "Create", "+" buttons
- If the goal is already accomplished, return action="done"
- Consider modern UI patterns (modals, dropdowns, contenteditable divs)
- If you need to navigate to a specific section (like Projects, Issues, etc), PREFER using action="navigate" with the direct URL instead of clicking sidebar links
- Look at the current URL to infer patterns (e.g., if current is linear.app/team/issues, Projects would be linear.app/team/projects)"""

            # Call GPT-4V
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": prompt
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/png;base64,{screenshot_base64}",
                                    "detail": "high"
                                }
                            }
                        ]
                    }
                ],
                max_tokens=500
            )

            result_text = response.choices[0].message.content.strip()

            # Clean JSON if wrapped in markdown
            if "```json" in result_text:
                result_text = result_text.split("```json")[1].split("```")[0]
            elif "```" in result_text:
                result_text = result_text.split("```")[1].split("```")[0]

            next_action = json.loads(result_text.strip())

            logger.info(f"GPT-4V next action: {next_action['action']} - {next_action.get('target', next_action.get('url', 'N/A'))}")
            logger.info(f"Reasoning: {next_action['reasoning']}")
            logger.info(f"Confidence: {next_action['confidence']}")

            return next_action

        except Exception as e:
            logger.error(f"Error getting next action from GPT-4V: {e}")
            # Return a wait action as fallback
            return {
                "action": "done",
                "reasoning": f"Error analyzing screenshot: {str(e)}",
                "confidence": 0.0
            }

    def verify_action_result(
        self,
        screenshot_path: str,
        action_taken: Dict[str, Any],
        expected_result: str
    ) -> Dict[str, Any]:
        """
        Verify if an action was successful by analyzing the new screenshot.

        Args:
            screenshot_path: Path to screenshot after action
            action_taken: The action that was executed
            expected_result: What we expected to happen

        Returns:
            Verification result with success, confidence, feedback
        """
        try:
            # Read and encode screenshot
            with open(screenshot_path, 'rb') as f:
                screenshot_bytes = f.read()
            screenshot_base64 = base64.b64encode(screenshot_bytes).decode('utf-8')

            prompt = f"""Verify if the following action was successful by analyzing the screenshot.

Action taken: {action_taken['action']} - {action_taken.get('target', action_taken.get('url', 'N/A'))}
Expected result: {expected_result}

Analyze the screenshot and determine:
1. Was the action successful?
2. How confident are you?
3. What evidence supports your conclusion?

Return ONLY a JSON object:
{{
  "success": true/false,
  "confidence": 0.0-1.0,
  "evidence": "what you see in the screenshot that indicates success/failure",
  "suggestion": "if failed, suggest what to try next (optional)"
}}"""

            # Call GPT-4V
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": prompt
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/png;base64,{screenshot_base64}",
                                    "detail": "high"
                                }
                            }
                        ]
                    }
                ],
                max_tokens=300
            )

            result_text = response.choices[0].message.content.strip()

            # Clean JSON if wrapped in markdown
            if "```json" in result_text:
                result_text = result_text.split("```json")[1].split("```")[0]
            elif "```" in result_text:
                result_text = result_text.split("```")[1].split("```")[0]

            verification = json.loads(result_text.strip())
            logger.info(f"Verification: {verification['success']} (confidence: {verification['confidence']})")
            logger.info(f"Evidence: {verification['evidence']}")

            return verification

        except Exception as e:
            logger.error(f"Error verifying action: {e}")
            return {
                "success": False,
                "confidence": 0.0,
                "evidence": f"Error: {str(e)}",
                "suggestion": "Retry the action"
            }