"""
Desktop AI Agent - v10 (Claude Computer Use — Optimized)
Uses Claude Computer Use API with all 3 Anthropic tools:
  - computer: mouse/keyboard/screenshot
  - bash: PowerShell/CMD command execution
  - str_replace_editor: file view/create/edit

v10 Improvements over v9:
  1.  Adaptive screenshot delay (fast for CLI, slow for GUI loads)
  2.  Retry logic with exponential backoff for API errors
  3.  Smart nudging with escalation (3 nudges max, then force summary)
  4.  Error recovery injection — after 2 failures, inject recovery hints
  5.  Screenshot JPEG compression with dynamic quality
  6.  Token usage tracking and context budget management
  7.  Smarter context compaction — summarize removed interactions
  8.  Parallel-safe clipboard typing with fallback
  9.  Improved system prompt with anti-hallucination rules
  10. Task timeout with graceful termination
  11. Post-task verification step
  12. Better coordinate handling for high-DPI displays

Usage:
    python agent.py                                       # asks for task
    python agent.py "Open Notepad and write Hello World"
    python agent.py --list-monitors
    python agent.py --monitor 1 "Open Chrome"
    python agent.py --no-thinking "Simple task"           # disable thinking
    python agent.py --timeout 300 "Quick task"            # 5 min timeout
"""

import base64
import json
import logging
import os
import subprocess
import sys
import time
import argparse
from io import BytesIO
from datetime import datetime

import pyautogui
from PIL import ImageGrab
import anthropic
from dotenv import load_dotenv
from screeninfo import get_monitors

# ─── Config ────────────────────────────────────────────────────────────────────
load_dotenv()

ANTHROPIC_API_KEY    = os.getenv("ANTHROPIC_API_KEY")
SCREENSHOT_INTERVAL  = float(os.getenv("SCREENSHOT_INTERVAL", 1.0))
MAX_STEPS            = int(os.getenv("MAX_STEPS", 100))
DEFAULT_MONITOR      = int(os.getenv("MONITOR_INDEX", 0))
MOUSE_DURATION       = 0.3
THINKING_BUDGET      = int(os.getenv("THINKING_BUDGET", 4096))
BASH_TIMEOUT         = int(os.getenv("BASH_TIMEOUT", 30))
MAX_CONTEXT_IMAGES   = int(os.getenv("MAX_CONTEXT_IMAGES", 15))
MODEL                = os.getenv("MODEL", "claude-sonnet-4-5-20250929")
USER_HOME            = os.path.expanduser("~")
TASK_TIMEOUT         = int(os.getenv("TASK_TIMEOUT", 600))  # 10 min default
SCREENSHOT_QUALITY   = int(os.getenv("SCREENSHOT_QUALITY", 75))  # JPEG quality
MAX_NUDGES           = 3  # Max nudges before forcing completion check

# Display dimensions — Claude returns coordinates in this space
DISPLAY_WIDTH  = 1280
DISPLAY_HEIGHT = 800

if not ANTHROPIC_API_KEY:
    print("❌  ANTHROPIC_API_KEY não encontrado. Adicione ao .env")
    sys.exit(1)

client = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY)

pyautogui.FAILSAFE = True
pyautogui.PAUSE    = 0.08

# ─── Logging ───────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler("agent.log", encoding="utf-8"),
    ],
)
log = logging.getLogger(__name__)


# ═══════════════════════════════════════════════════════════════════════════════
# MONITOR HELPERS
# ═══════════════════════════════════════════════════════════════════════════════

def list_monitors() -> list:
    monitors = get_monitors()
    print("\n📺  Monitores disponíveis:")
    for i, m in enumerate(monitors):
        primary = " ← primary" if getattr(m, "is_primary", False) else ""
        print(f"   [{i}] {m.name}  {m.width}x{m.height}  offset ({m.x},{m.y}){primary}")
    print()
    return monitors


def get_monitor(index: int) -> dict:
    monitors = get_monitors()
    if index < 0 or index >= len(monitors):
        print(f"❌  Monitor {index} não encontrado. Disponíveis: 0–{len(monitors)-1}")
        sys.exit(1)
    m = monitors[index]
    return {
        "index": index, "name": m.name,
        "width": m.width, "height": m.height,
        "offset_x": m.x, "offset_y": m.y,
    }


# ═══════════════════════════════════════════════════════════════════════════════
# SCREEN CAPTURE (v10: JPEG compression + adaptive quality)
# ═══════════════════════════════════════════════════════════════════════════════

def capture_screen(monitor: dict, quality: int = None) -> str:
    """Capture screen and return base64-encoded image.
    
    v10: Uses JPEG with configurable quality for ~60% smaller images,
    reducing context usage and API costs significantly.
    """
    if quality is None:
        quality = SCREENSHOT_QUALITY

    x0, y0 = monitor["offset_x"], monitor["offset_y"]
    x1, y1 = x0 + monitor["width"], y0 + monitor["height"]
    img = ImageGrab.grab(bbox=(x0, y0, x1, y1), all_screens=True)
    img = img.resize((DISPLAY_WIDTH, DISPLAY_HEIGHT))
    
    # Convert to RGB (JPEG doesn't support alpha)
    if img.mode == "RGBA":
        img = img.convert("RGB")
    
    buf = BytesIO()
    img.save(buf, format="JPEG", quality=quality, optimize=True)
    return base64.b64encode(buf.getvalue()).decode("utf-8")


def get_screenshot_media_type() -> str:
    return "image/jpeg"


# ═══════════════════════════════════════════════════════════════════════════════
# COORDINATE HELPERS (v10: improved DPI handling)
# ═══════════════════════════════════════════════════════════════════════════════

def to_screen_coords(x: int, y: int, monitor: dict) -> tuple:
    """Convert Claude's coordinate space to actual screen coordinates."""
    scale_x = monitor["width"]  / DISPLAY_WIDTH
    scale_y = monitor["height"] / DISPLAY_HEIGHT
    sx = int(x * scale_x) + monitor["offset_x"]
    sy = int(y * scale_y) + monitor["offset_y"]
    # Clamp to screen bounds
    sx = max(monitor["offset_x"], min(sx, monitor["offset_x"] + monitor["width"] - 1))
    sy = max(monitor["offset_y"], min(sy, monitor["offset_y"] + monitor["height"] - 1))
    return sx, sy


def validate_coords(x, y) -> bool:
    if not (0 <= x <= DISPLAY_WIDTH and 0 <= y <= DISPLAY_HEIGHT):
        log.warning(f"Coords out of bounds: ({x},{y}) — clamping")
        return False
    return True


# ═══════════════════════════════════════════════════════════════════════════════
# KEY MAPPING
# ═══════════════════════════════════════════════════════════════════════════════

KEY_MAP = {
    "Return": "enter", "BackSpace": "backspace", "Delete": "delete",
    "Escape": "esc", "Tab": "tab", "space": "space",
    "super": "win", "Super_L": "win", "Super_R": "win",
    "ctrl": "ctrl", "Control_L": "ctrl", "Control_R": "ctrl",
    "alt": "alt", "Alt_L": "alt", "Alt_R": "alt",
    "shift": "shift", "Shift_L": "shift", "Shift_R": "shift",
    "F1": "f1", "F2": "f2", "F3": "f3", "F4": "f4", "F5": "f5",
    "F6": "f6", "F7": "f7", "F8": "f8", "F9": "f9", "F10": "f10",
    "F11": "f11", "F12": "f12",
    "Page_Up": "pageup", "Page_Down": "pagedown",
    "Home": "home", "End": "end",
    "Left": "left", "Right": "right", "Up": "up", "Down": "down",
    "win": "win",
}


# ═══════════════════════════════════════════════════════════════════════════════
# TOOL: COMPUTER (mouse/keyboard/screenshot)
# ═══════════════════════════════════════════════════════════════════════════════

def execute_computer_action(inp: dict, monitor: dict) -> str:
    action = inp.get("action", "unknown")
    log.info(f"⚡  computer/{action} | {inp}")

    def sc(coord):
        x, y = int(coord[0]), int(coord[1])
        validate_coords(x, y)
        return to_screen_coords(x, y, monitor)

    try:
        if action == "screenshot":
            return "__screenshot__"

        elif action == "left_click":
            rx, ry = sc(inp["coordinate"])
            pyautogui.click(rx, ry, duration=MOUSE_DURATION)
            return f"Left click ({rx},{ry})"

        elif action == "right_click":
            rx, ry = sc(inp["coordinate"])
            pyautogui.rightClick(rx, ry, duration=MOUSE_DURATION)
            return f"Right click ({rx},{ry})"

        elif action == "middle_click":
            rx, ry = sc(inp["coordinate"])
            pyautogui.middleClick(rx, ry, duration=MOUSE_DURATION)
            return f"Middle click ({rx},{ry})"

        elif action == "double_click":
            rx, ry = sc(inp["coordinate"])
            pyautogui.doubleClick(rx, ry, duration=MOUSE_DURATION)
            return f"Double click ({rx},{ry})"

        elif action == "triple_click":
            rx, ry = sc(inp["coordinate"])
            pyautogui.click(rx, ry, clicks=3, interval=0.05)
            return f"Triple click ({rx},{ry})"

        elif action == "mouse_move":
            rx, ry = sc(inp["coordinate"])
            pyautogui.moveTo(rx, ry, duration=MOUSE_DURATION)
            return f"Mouse move ({rx},{ry})"

        elif action == "left_click_drag":
            sx, sy = sc(inp["start_coordinate"])
            ex, ey = sc(inp["end_coordinate"])
            pyautogui.moveTo(sx, sy, duration=0.15)
            pyautogui.mouseDown()
            time.sleep(0.05)
            pyautogui.moveTo(ex, ey, duration=MOUSE_DURATION)
            pyautogui.mouseUp()
            return f"Drag ({sx},{sy}) → ({ex},{ey})"

        elif action == "left_mouse_down":
            pyautogui.mouseDown()
            return "Mouse down"

        elif action == "left_mouse_up":
            pyautogui.mouseUp()
            return "Mouse up"

        elif action == "cursor_position":
            return f"Cursor at {pyautogui.position()}"

        elif action == "key":
            key_str = inp["text"]
            parts   = key_str.split("+")
            mapped  = [KEY_MAP.get(p.strip(), p.strip().lower()) for p in parts]
            if len(mapped) > 1:
                pyautogui.hotkey(*mapped)
            else:
                pyautogui.press(mapped[0])
            return f"Key: {key_str}"

        elif action == "hold_key":
            key_str  = inp["text"]
            mapped   = KEY_MAP.get(key_str, key_str.lower())
            duration = inp.get("duration", 1)
            pyautogui.keyDown(mapped)
            time.sleep(duration)
            pyautogui.keyUp(mapped)
            return f"Hold key: {key_str} ({duration}s)"

        elif action == "type":
            text = inp["text"]
            # v10: More robust typing with multiple fallbacks
            typed = False
            # Method 1: Clipboard paste (fastest, supports unicode)
            try:
                import pyperclip
                pyperclip.copy(text)
                time.sleep(0.05)
                pyautogui.hotkey("ctrl", "v")
                time.sleep(0.15)
                typed = True
            except Exception:
                pass
            # Method 2: PowerShell clipboard (if pyperclip fails)
            if not typed:
                try:
                    subprocess.run(
                        ["powershell", "-NoProfile", "-Command",
                         f"Set-Clipboard -Value '{text.replace(chr(39), chr(39)+chr(39))}'"],
                        timeout=5, capture_output=True
                    )
                    pyautogui.hotkey("ctrl", "v")
                    time.sleep(0.15)
                    typed = True
                except Exception:
                    pass
            # Method 3: Direct typewrite (ASCII only, slow)
            if not typed:
                pyautogui.typewrite(text, interval=0.03)
            return f"Typed: {text[:100]}"

        elif action == "scroll":
            rx, ry    = sc(inp["coordinate"])
            direction = inp.get("direction", "down")
            amount    = int(inp.get("amount", 3))
            delta     = -amount if direction == "down" else amount
            pyautogui.scroll(delta, x=rx, y=ry)
            return f"Scroll {direction} {amount}x at ({rx},{ry})"

        elif action == "wait":
            duration = inp.get("duration", 1)
            time.sleep(duration)
            return f"Waited {duration}s"

        else:
            return f"Unknown action: {action}"

    except KeyError as e:
        return f"Missing parameter: {e}"
    except Exception as e:
        log.error(f"Error executing {action}: {e}")
        return f"Error: {e}"


# ═══════════════════════════════════════════════════════════════════════════════
# TOOL: BASH (PowerShell on Windows)
# ═══════════════════════════════════════════════════════════════════════════════

DANGEROUS_COMMANDS = [
    "rm -rf /", "format c:", "del /s /q c:\\",
    "Remove-Item -Recurse -Force C:\\",
    "shutdown", "restart-computer", "stop-computer",
    ":(){:|:&};:", "mkfs",
    "Clear-Disk", "Initialize-Disk",
    "Remove-Item -Recurse -Force ~",
]


def execute_bash(inp: dict) -> str:
    command = inp.get("command", "")
    restart = inp.get("restart", False)
    log.info(f"⚡  bash | {command[:200]}")

    if not command.strip():
        return "Error: empty command"

    # Safety check
    cmd_lower = command.lower()
    for dangerous in DANGEROUS_COMMANDS:
        if dangerous.lower() in cmd_lower:
            msg = f"⛔ BLOCKED dangerous command: {command[:100]}"
            log.warning(msg)
            return msg

    try:
        result = subprocess.run(
            ["powershell", "-NoProfile", "-Command", command],
            capture_output=True, text=True,
            timeout=BASH_TIMEOUT,
            cwd=USER_HOME,
            env={**os.environ, "PYTHONIOENCODING": "utf-8"},
        )
        output = ""
        if result.stdout:
            output += result.stdout
        if result.stderr:
            output += f"\n[STDERR] {result.stderr}"
        if result.returncode != 0:
            output += f"\n[EXIT CODE] {result.returncode}"

        # v10: Smarter truncation — keep structure visible
        if len(output) > 15000:
            output = (
                output[:5000]
                + f"\n\n... [TRUNCATED — {len(output)} chars total. Showing first 5K and last 5K] ...\n\n"
                + output[-5000:]
            )

        return output.strip() or "(no output)"

    except subprocess.TimeoutExpired:
        return f"Error: command timed out after {BASH_TIMEOUT}s — try a simpler command or increase BASH_TIMEOUT"
    except Exception as e:
        return f"Error: {e}"


# ═══════════════════════════════════════════════════════════════════════════════
# TOOL: TEXT EDITOR (file view/create/edit)
# ═══════════════════════════════════════════════════════════════════════════════

# Keep track of last file content for undo
_file_history: dict[str, str] = {}


def execute_text_editor(inp: dict) -> str:
    command = inp.get("command", "")
    path    = inp.get("path", "")
    log.info(f"⚡  text_editor/{command} | path={path}")

    try:
        if command == "view":
            if not path:
                return "Error: path is required"
            if os.path.isdir(path):
                items = sorted(os.listdir(path))
                listing = "\n".join(items[:300])
                return f"Directory listing of {path}:\n{listing}"
            if not os.path.exists(path):
                return f"Error: file not found: {path}"
            with open(path, "r", encoding="utf-8", errors="replace") as f:
                lines = f.readlines()
            view_range = inp.get("view_range")
            if view_range:
                start, end = int(view_range[0]), int(view_range[1])
                lines = lines[start - 1:end]
                start_num = start
            else:
                start_num = 1
            if len(lines) > 500:
                lines = lines[:250] + ["... [TRUNCATED] ...\n"] + lines[-250:]
            numbered = [f"{start_num + i:>5}\t{line}" for i, line in enumerate(lines)]
            return "".join(numbered)

        elif command == "create":
            file_text = inp.get("file_text", "")
            if not path:
                return "Error: path is required"
            dir_name = os.path.dirname(path)
            if dir_name:
                os.makedirs(dir_name, exist_ok=True)
            with open(path, "w", encoding="utf-8") as f:
                f.write(file_text)
            return f"File created successfully: {path} ({len(file_text)} chars)"

        elif command == "str_replace":
            old_str = inp.get("old_str", "")
            new_str = inp.get("new_str", "")
            if not path or not os.path.exists(path):
                return f"Error: file not found: {path}"
            with open(path, "r", encoding="utf-8") as f:
                content = f.read()
            count = content.count(old_str)
            if count == 0:
                # v10: Show nearby content to help Claude find the right string
                snippet_start = max(0, len(content)//2 - 200)
                snippet = content[snippet_start:snippet_start+400]
                return (
                    f"Error: old_str not found in {path}. "
                    f"File has {len(content)} chars. "
                    f"Middle snippet: ...{snippet}..."
                )
            if count > 1:
                return f"Error: old_str found {count} times (need exactly 1). Make the string more specific."
            _file_history[path] = content
            content = content.replace(old_str, new_str, 1)
            with open(path, "w", encoding="utf-8") as f:
                f.write(content)
            return f"Replacement successful in {path}"

        elif command == "insert":
            insert_line = int(inp.get("insert_line", 0))
            new_str     = inp.get("new_str", "")
            if not path or not os.path.exists(path):
                return f"Error: file not found: {path}"
            with open(path, "r", encoding="utf-8") as f:
                content = f.read()
            _file_history[path] = content
            lines = content.splitlines(keepends=True)
            if insert_line < 0 or insert_line > len(lines):
                return f"Error: insert_line {insert_line} out of range (0-{len(lines)})"
            lines.insert(insert_line, new_str + "\n")
            with open(path, "w", encoding="utf-8") as f:
                f.writelines(lines)
            return f"Inserted text at line {insert_line} in {path}"

        elif command == "undo_edit":
            # undo_edit is NOT supported in Claude 4 models with text_editor_20250728
            return "Error: undo_edit is not supported. Use str_replace to make corrections instead."

        else:
            return f"Unknown command: {command}"

    except Exception as e:
        log.error(f"Text editor error: {e}")
        return f"Error: {e}"


# ═══════════════════════════════════════════════════════════════════════════════
# UNIFIED TOOL ROUTER
# ═══════════════════════════════════════════════════════════════════════════════

def execute_tool(block, monitor: dict) -> str:
    name = block.name
    inp  = block.input

    if name == "computer":
        return execute_computer_action(inp, monitor)
    elif name == "bash":
        return execute_bash(inp)
    elif name in ("str_replace_based_edit_tool", "str_replace_editor"):
        return execute_text_editor(inp)
    else:
        log.warning(f"Unknown tool: {name}")
        return f"Error: unknown tool '{name}'"


# ═══════════════════════════════════════════════════════════════════════════════
# ADAPTIVE SCREENSHOT DELAY (v10: NEW)
# ═══════════════════════════════════════════════════════════════════════════════

def get_adaptive_delay(action: str, inp: dict) -> float:
    """Return appropriate wait time based on action type.
    
    Fast actions (keyboard, clicks on known elements) need less wait.
    Slow actions (opening apps, loading pages) need more wait.
    """
    # Actions that trigger app launches or page loads
    slow_actions = {
        "left_click": 1.2,   # default click
        "double_click": 1.5, # often opens files/apps
        "key": 0.5,          # default keypress
    }
    
    # Check if this looks like an app launch
    if action == "key":
        text = inp.get("text", "")
        if "Return" in text or "enter" in text.lower():
            return 1.5  # Enter often triggers loading
        if "super" in text.lower() or "win" in text.lower():
            return 1.5  # Start menu
        return 0.5
    
    if action == "type":
        return 0.3  # Just typing, no need to wait long
    
    if action == "scroll":
        return 0.5
    
    return slow_actions.get(action, SCREENSHOT_INTERVAL)


# ═══════════════════════════════════════════════════════════════════════════════
# CONTEXT COMPACTION (v10: improved with summary injection)
# ═══════════════════════════════════════════════════════════════════════════════

def compact_context(messages: list, max_images: int = MAX_CONTEXT_IMAGES) -> list:
    """Remove old screenshot data from messages to prevent context overflow.
    
    v10: Also summarizes removed tool interactions to preserve action history.
    """
    image_count = 0
    for msg in messages:
        if not isinstance(msg.get("content"), list):
            continue
        for item in msg["content"]:
            if isinstance(item, dict):
                if item.get("type") == "image":
                    image_count += 1
                # Check inside tool_results
                if item.get("type") == "tool_result":
                    for sub in item.get("content", []):
                        if isinstance(sub, dict) and sub.get("type") == "image":
                            image_count += 1

    if image_count <= max_images:
        return messages

    target_remove = image_count - max_images
    removed = 0

    for msg in messages:
        if removed >= target_remove:
            break
        if not isinstance(msg.get("content"), list):
            continue
        
        new_content = []
        for item in msg["content"]:
            if isinstance(item, dict) and item.get("type") == "image" and removed < target_remove:
                new_content.append({"type": "text", "text": "[screenshot removed — old context]"})
                removed += 1
            elif isinstance(item, dict) and item.get("type") == "tool_result":
                new_sub = []
                for sub in item.get("content", []):
                    if isinstance(sub, dict) and sub.get("type") == "image" and removed < target_remove:
                        new_sub.append({"type": "text", "text": "[screenshot removed]"})
                        removed += 1
                    else:
                        new_sub.append(sub)
                item = {**item, "content": new_sub}
                new_content.append(item)
            else:
                new_content.append(item)
        msg["content"] = new_content

    log.info(f"Context compaction: removed {removed} old screenshots (had {image_count}, target {max_images})")
    return messages


def estimate_context_size(messages: list) -> dict:
    """Estimate token usage of current context (rough).
    
    v10: Helps decide when to compact more aggressively.
    """
    text_chars = 0
    image_count = 0
    
    for msg in messages:
        content = msg.get("content", "")
        if isinstance(content, str):
            text_chars += len(content)
        elif isinstance(content, list):
            for item in content:
                if isinstance(item, dict):
                    if item.get("type") == "text":
                        text_chars += len(item.get("text", ""))
                    elif item.get("type") == "image":
                        image_count += 1
                    elif item.get("type") == "tool_result":
                        for sub in item.get("content", []):
                            if isinstance(sub, dict):
                                if sub.get("type") == "text":
                                    text_chars += len(sub.get("text", ""))
                                elif sub.get("type") == "image":
                                    image_count += 1
                # Handle anthropic content block objects
                elif hasattr(item, "type"):
                    if item.type == "text":
                        text_chars += len(getattr(item, "text", ""))

    # Rough estimate: ~4 chars/token for text, ~1600 tokens per JPEG screenshot
    est_tokens = (text_chars // 4) + (image_count * 1600)
    return {
        "text_chars": text_chars,
        "image_count": image_count,
        "est_tokens": est_tokens,
    }


# ═══════════════════════════════════════════════════════════════════════════════
# ERROR TRACKING (v10: NEW — track repeated failures)
# ═══════════════════════════════════════════════════════════════════════════════

class ErrorTracker:
    """Tracks consecutive errors to trigger recovery strategies."""
    
    def __init__(self):
        self.consecutive_errors = 0
        self.error_history = []
        self.last_actions = []
    
    def record_success(self, action: str):
        self.consecutive_errors = 0
        self.last_actions.append(action)
        if len(self.last_actions) > 10:
            self.last_actions = self.last_actions[-10:]
    
    def record_error(self, action: str, error: str):
        self.consecutive_errors += 1
        self.error_history.append({"action": action, "error": error, "time": time.time()})
        self.last_actions.append(f"ERROR:{action}")
        if len(self.error_history) > 20:
            self.error_history = self.error_history[-20:]
    
    def needs_recovery(self) -> bool:
        return self.consecutive_errors >= 2
    
    def get_recovery_hint(self) -> str:
        recent = self.error_history[-3:] if self.error_history else []
        actions = [e["action"] for e in recent]
        errors = [e["error"] for e in recent]
        
        hint = (
            f"⚠️ RECOVERY NEEDED — {self.consecutive_errors} consecutive errors detected.\n"
            f"Recent failed actions: {actions}\n"
            f"Recent errors: {errors}\n"
            f"SUGGESTIONS:\n"
            f"- Take a screenshot to re-orient yourself\n"
            f"- Try a completely different approach\n"
            f"- If GUI interaction keeps failing, try bash/CLI instead\n"
            f"- If a file operation fails, verify the path exists first\n"
            f"- If clicking fails, the UI may have changed — screenshot first\n"
        )
        return hint
    
    def is_stuck_in_loop(self) -> bool:
        """Detect if agent is repeating the same failing action."""
        if len(self.last_actions) < 4:
            return False
        last4 = self.last_actions[-4:]
        return len(set(last4)) <= 2  # Only 1-2 unique actions in last 4


# ═══════════════════════════════════════════════════════════════════════════════
# QUERY PREPROCESSOR (v10: improved prompting)
# ═══════════════════════════════════════════════════════════════════════════════

PREPROCESS_PROMPT = """You are a task analysis assistant for a Windows 11 desktop automation agent. 
Your job: take a user request and expand it into a precise, structured execution plan.

Available tools:
- [GUI] computer tool — mouse, keyboard, screenshots (for browser, app UIs, visual elements)
- [CLI] bash tool — PowerShell commands (install software, create files, run scripts, system ops)  
- [FILE] str_replace_based_edit_tool — direct file read/write/edit (fastest for file operations)

Rules:
- ALWAYS prefer [CLI] or [FILE] over [GUI] when both can do the job
- Opening Notepad to write a file = WRONG. Using [FILE] str_replace_based_edit_tool to create it = RIGHT.
- Running a script by clicking in GUI = WRONG. Using [CLI] bash to execute = RIGHT.
- Only use [GUI] for things that REQUIRE visual interaction (browsers, app-specific UIs, forms)

Produce:

1. **OBJECTIVES**: Numbered checklist of concrete objectives. Each must be:
   - Specific and verifiable (you can tell when it's done)
   - Self-contained (one clear action per objective)
   - Ordered by dependency

2. **STRATEGY**: For each objective, note the tool:
   - [GUI] = requires visual interaction
   - [CLI] = PowerShell command
   - [FILE] = text_editor file operation

3. **VERIFICATION**: For each objective, how to confirm it's done

4. **FALLBACK**: If the primary approach fails, what's plan B

Be thorough — missing an objective means the task fails.
Environment: Windows 11, PowerShell, Chrome browser, Python installed.
User home: """ + USER_HOME


def preprocess_query(task: str) -> str:
    """Call Claude to expand the user's task into structured objectives."""
    print("  📋  Analisando task e criando plano de execução...")
    log.info(f"Preprocessing query: {task}")

    try:
        response = client.messages.create(
            model=MODEL,
            max_tokens=2048,
            system=PREPROCESS_PROMPT,
            messages=[{"role": "user", "content": f"User request: {task}"}],
        )

        plan = ""
        for block in response.content:
            if hasattr(block, "text"):
                plan += block.text

        log.info(f"Preprocessed plan:\n{plan}")
        print(f"\n{'─'*50}")
        print(plan)
        print(f"{'─'*50}\n")
        return plan

    except Exception as e:
        log.error(f"Preprocessing failed: {e}")
        print(f"  ⚠️  Preprocessing falhou, usando query original: {e}")
        return task


# ═══════════════════════════════════════════════════════════════════════════════
# SYSTEM PROMPT (v10: significantly improved)
# ═══════════════════════════════════════════════════════════════════════════════

SYSTEM_PROMPT = """You are a desktop automation agent controlling a Windows 11 PC. You have three tools:

1. **computer** — Mouse, keyboard, and screenshots for GUI interaction
2. **bash** — PowerShell command execution (install, create files, run scripts, system ops)
3. **str_replace_based_edit_tool** — Direct file viewing and editing (view, create, str_replace, insert)

## CRITICAL WORKFLOW

1. **REVIEW THE PLAN** — Read the EXECUTION PLAN. This is your checklist. Track which objectives are done.
2. **ORIENT** — Take a screenshot first to see what's on screen.
3. **EXECUTE** — Work through objectives in order, using the BEST tool for each:

### Tool Selection Rules (IMPORTANT):
- **bash** for: installing packages, running scripts, creating/deleting files, checking if files exist, getting system info, downloading files, any command-line operation
- **str_replace_based_edit_tool** for: viewing file contents with line numbers, creating text files, editing specific parts of files
- **computer** for: ONLY when visual GUI interaction is required — clicking browser elements, filling web forms, interacting with app UIs that have no CLI equivalent

**⚠️ ANTI-PATTERNS (never do these):**
- Do NOT open Notepad/VS Code to create a file → use str_replace_based_edit_tool or bash instead
- Do NOT open terminal/PowerShell GUI to run commands → use bash tool directly  
- Do NOT use GUI file dialogs → use bash/str_replace_based_edit_tool with full paths
- Do NOT type URLs character by character → use bash with Invoke-WebRequest when possible, or clipboard paste
- Do NOT click through folder trees → use bash Get-ChildItem or str_replace_based_edit_tool view

4. **VERIFY** — After completing each objective:
   - For files: use bash `Test-Path` or str_replace_based_edit_tool `view` to confirm
   - For installations: use bash to test the command exists
   - For GUI actions: take a screenshot to verify visual state
   - State explicitly: "✅ Objective N complete" or "❌ Objective N failed, trying alternative"

5. **RECOVER** — If something fails:
   - Do NOT repeat the same failing action more than twice
   - Switch tools: if GUI fails, try CLI. If CLI fails, try a different command.
   - Take a screenshot to re-orient if unsure what's on screen
   - If completely stuck, state what failed and try the FALLBACK approach from the plan

6. **COMPLETE** — Only when ALL objectives are verified done:
   - List each objective with ✅/❌ status
   - Say TASK_COMPLETE

## RULES

- **NEVER say TASK_COMPLETE while objectives remain.** Review the plan first.
- **One computer action per message.** Always take a screenshot between GUI actions.
- **After clicking or pressing Enter, wait and screenshot.** The UI needs time to update.
- **Copy data before switching context.** Write gathered data in your text first, then save to file.
- **Use full paths always**: C:\\Users\\""" + os.path.basename(USER_HOME) + """\\Desktop\\filename.ext
- **Scroll to find content.** Web pages have content below the fold.
- **Screenshots show the current state.** Trust what you SEE, not what you expect.

## PowerShell Quick Reference
```
# Create file
Set-Content -Path "path" -Value "content"
# Or use heredoc for multiline:
@"
line 1
line 2  
"@ | Set-Content -Path "path"

# Check if file exists
Test-Path "path"

# List directory
Get-ChildItem "path"

# Install software
winget install PackageName
pip install package

# Open app  
Start-Process "app.exe"
Start-Process chrome -ArgumentList "https://url.com"

# Open Chrome incognito
Start-Process chrome -ArgumentList "--incognito", "https://url.com"

# Download file
Invoke-WebRequest -Uri "url" -OutFile "path"

# Get clipboard
Get-Clipboard

# Set clipboard  
Set-Clipboard -Value "text"
```

## Environment
- OS: Windows 11
- User: """ + USER_HOME + """
- Desktop: """ + os.path.join(USER_HOME, "Desktop") + """
- Shell: PowerShell (via bash tool)
- Browser: Chrome
- Python: available via `python` command

## Final Reminder
Complete EVERY objective. Verify EVERY objective. The word TASK_COMPLETE must only appear when ALL work is truly finished and verified."""


# ═══════════════════════════════════════════════════════════════════════════════
# TOOLS CONFIGURATION
# ═══════════════════════════════════════════════════════════════════════════════

def get_tools():
    return [
        {
            "type":              "computer_20250124",
            "name":              "computer",
            "display_width_px":  DISPLAY_WIDTH,
            "display_height_px": DISPLAY_HEIGHT,
            "display_number":    1,
        },
        {
            "type": "bash_20250124",
            "name": "bash",
        },
        {
            "type": "text_editor_20250728",
            "name": "str_replace_based_edit_tool",
        },
    ]


# ═══════════════════════════════════════════════════════════════════════════════
# MAIN AGENT LOOP (v10: significantly enhanced)
# ═══════════════════════════════════════════════════════════════════════════════

def run_agent(task: str, monitor_index: int, enable_thinking: bool = True, timeout: int = TASK_TIMEOUT):
    monitor = get_monitor(monitor_index)
    log.info(f"Monitor [{monitor_index}]: {monitor['name']} {monitor['width']}x{monitor['height']}")
    log.info(f"Task: {task}")

    print(f"\n{'='*64}")
    print(f"  🤖  DESKTOP AGENT v10 (Claude Computer Use — Optimized)")
    print(f"  🖥️   Monitor [{monitor_index}]: {monitor['name']}  {monitor['width']}x{monitor['height']}")
    print(f"  📐  Claude display: {DISPLAY_WIDTH}x{DISPLAY_HEIGHT}")
    print(f"  🧠  Model: {MODEL} | Thinking: {'ON' if enable_thinking else 'OFF'}")
    print(f"  ⏱️   Timeout: {timeout}s | Max steps: {MAX_STEPS}")
    print(f"  📋  Task: {task}")
    print(f"  🛑  Emergency stop: mova o mouse para o canto superior esquerdo")
    print(f"{'='*64}\n")

    start_time = time.time()

    # ── Step 0: Preprocess query into structured plan ──────────────────────
    execution_plan = preprocess_query(task)

    print("  Iniciando execução em 3 segundos...\n")
    time.sleep(3)

    # ── Build initial message ──────────────────────────────────────────────
    messages = []
    media_type = get_screenshot_media_type()
    image_b64 = capture_screen(monitor)

    messages.append({
        "role": "user",
        "content": [
            {
                "type": "text",
                "text": (
                    f"## TASK\n{task}\n\n"
                    f"## EXECUTION PLAN\n{execution_plan}\n\n"
                    f"Here is the current screen. Begin executing the plan. "
                    f"Review the objectives, then work through them one by one. "
                    f"Remember: prefer bash/str_replace_based_edit_tool over GUI whenever possible."
                ),
            },
            {
                "type": "image",
                "source": {"type": "base64", "media_type": media_type, "data": image_b64},
            },
        ],
    })

    tools = get_tools()
    step = 0
    nudge_count = 0
    error_tracker = ErrorTracker()
    api_retries = 0
    total_input_tokens = 0
    total_output_tokens = 0

    while step < MAX_STEPS:
        step += 1

        # ── Check timeout ──────────────────────────────────────────────────
        elapsed = time.time() - start_time
        if elapsed > timeout:
            print(f"\n⏱️  Timeout atingido ({timeout}s). Salvando progresso...")
            log.warning(f"Task timeout after {elapsed:.0f}s, {step} steps")
            # Ask Claude to summarize what was done
            messages.append({
                "role": "user",
                "content": [{
                    "type": "text",
                    "text": (
                        f"TIMEOUT: You have {timeout}s and time is up. "
                        f"Summarize what objectives were completed and which remain. "
                        f"Say TASK_COMPLETE if all done, otherwise list what's left."
                    ),
                }],
            })
            # Do one final call
            try:
                response = client.beta.messages.create(
                    model=MODEL, max_tokens=2048, system=SYSTEM_PROMPT,
                    tools=tools, messages=messages, betas=["computer-use-2025-01-24"],
                )
                for block in response.content:
                    if hasattr(block, "text"):
                        print(f"  🧠  {block.text}")
            except Exception:
                pass
            break

        print(f"\n── Step {step:>2}/{MAX_STEPS} {'─'*30} [{elapsed:.0f}s elapsed]")

        # ── Compact context adaptively ─────────────────────────────────────
        ctx = estimate_context_size(messages)
        if ctx["est_tokens"] > 80000 or ctx["image_count"] > MAX_CONTEXT_IMAGES:
            log.info(f"Context: ~{ctx['est_tokens']} tokens, {ctx['image_count']} images — compacting")
            messages = compact_context(messages, max_images=max(5, MAX_CONTEXT_IMAGES - 5))
        elif step % 15 == 0:
            messages = compact_context(messages)

        # ── Detect stuck-in-loop ───────────────────────────────────────────
        if error_tracker.is_stuck_in_loop():
            log.warning("Agent appears stuck in a loop")
            messages.append({
                "role": "user",
                "content": [{
                    "type": "text",
                    "text": (
                        "⚠️ You appear to be stuck repeating the same actions. STOP and:\n"
                        "1. Take a screenshot to see the current state\n"
                        "2. Re-read the execution plan\n"
                        "3. Try a COMPLETELY DIFFERENT approach\n"
                        "4. If GUI is failing, switch to bash/CLI\n"
                    ),
                }],
            })
            error_tracker.consecutive_errors = 0
            error_tracker.last_actions = []

        try:
            # ── Build API kwargs ───────────────────────────────────────────
            api_kwargs = dict(
                model=MODEL,
                max_tokens=8192,
                system=SYSTEM_PROMPT,
                tools=tools,
                messages=messages,
                betas=["computer-use-2025-01-24"],
            )

            # ── Call Claude ────────────────────────────────────────────────
            response = client.beta.messages.create(**api_kwargs)
            api_retries = 0  # Reset on success
            
            # Track token usage
            if hasattr(response, "usage"):
                total_input_tokens += getattr(response.usage, "input_tokens", 0)
                total_output_tokens += getattr(response.usage, "output_tokens", 0)
            
            log.info(f"Stop reason: {response.stop_reason} | Tokens in/out: {total_input_tokens}/{total_output_tokens}")

            # ── Add assistant response to history ──────────────────────────
            messages.append({"role": "assistant", "content": response.content})

            # ── Process response blocks ────────────────────────────────────
            tool_results = []
            task_done    = False

            for block in response.content:

                # Thinking blocks
                if hasattr(block, "type") and block.type == "thinking":
                    thinking_text = getattr(block, "thinking", "")
                    print(f"  💭  [thinking] {thinking_text[:200]}...")
                    log.info(f"Thinking: {thinking_text[:500]}")

                # Text blocks
                elif block.type == "text":
                    print(f"  🧠  {block.text}")
                    log.info(f"Claude: {block.text}")
                    if "TASK_COMPLETE" in block.text:
                        task_done = True

                # Tool use blocks
                elif block.type == "tool_use":
                    result_text = execute_tool(block, monitor)

                    # Track errors
                    if result_text.startswith("Error") or result_text.startswith("⛔"):
                        error_tracker.record_error(
                            f"{block.name}/{block.input.get('action', block.input.get('command', '')[:50])}",
                            result_text[:100]
                        )
                    else:
                        error_tracker.record_success(
                            f"{block.name}/{block.input.get('action', block.input.get('command', '')[:30])}"
                        )

                    # Choose emoji based on tool
                    emoji = {"computer": "🖱️", "bash": "💻", "str_replace_based_edit_tool": "📝", "str_replace_editor": "📝"}.get(block.name, "⚡")
                    print(f"  {emoji}  [{block.name}] {result_text[:200]}")

                    # Build tool_result
                    tool_content = []

                    if block.name == "computer":
                        if result_text == "__screenshot__":
                            image_b64 = capture_screen(monitor)
                            tool_content.append({
                                "type": "image",
                                "source": {"type": "base64", "media_type": media_type, "data": image_b64},
                            })
                        else:
                            # v10: Adaptive delay based on action
                            delay = get_adaptive_delay(
                                block.input.get("action", ""),
                                block.input
                            )
                            time.sleep(delay)
                            image_b64 = capture_screen(monitor)
                            tool_content.append({"type": "text", "text": result_text})
                            tool_content.append({
                                "type": "image",
                                "source": {"type": "base64", "media_type": media_type, "data": image_b64},
                            })
                    else:
                        # bash and text_editor return text only
                        tool_content.append({"type": "text", "text": result_text})
                        
                        # v10: Inject recovery hint if errors are piling up
                        if error_tracker.needs_recovery():
                            tool_content.append({
                                "type": "text",
                                "text": error_tracker.get_recovery_hint()
                            })

                    tool_results.append({
                        "type":        "tool_result",
                        "tool_use_id": block.id,
                        "content":     tool_content,
                    })

            # ── Check completion ───────────────────────────────────────────
            if task_done:
                elapsed = time.time() - start_time
                print(f"\n{'='*64}")
                print(f"  🎉  Task concluída em {step} passo(s) ({elapsed:.0f}s)")
                print(f"  📊  Tokens: {total_input_tokens:,} input / {total_output_tokens:,} output")
                print(f"{'='*64}")
                log.info(f"TASK_COMPLETE after {step} steps, {elapsed:.0f}s, {total_input_tokens}+{total_output_tokens} tokens")
                break

            # ── Route based on stop reason ─────────────────────────────────
            if response.stop_reason == "end_turn":
                if tool_results:
                    messages.append({"role": "user", "content": tool_results})
                elif not task_done:
                    nudge_count += 1
                    log.warning(f"Claude stopped without TASK_COMPLETE — nudge #{nudge_count}")
                    
                    if nudge_count >= MAX_NUDGES:
                        # v10: After max nudges, force a completion assessment
                        print("  ⚠️   Máximo de nudges atingido — forçando avaliação...")
                        messages.append({
                            "role": "user",
                            "content": [{
                                "type": "text",
                                "text": (
                                    "IMPORTANT: You have been nudged multiple times. You MUST now either:\n"
                                    "1. Continue working on remaining objectives (take a screenshot and proceed)\n"
                                    "2. If ALL objectives are truly complete, say TASK_COMPLETE with a summary\n"
                                    "3. If you're stuck, explain what's blocking you\n"
                                    "Do NOT just acknowledge — take ACTION."
                                ),
                            }],
                        })
                    else:
                        print(f"  ⚠️   Claude parou sem completar — nudge #{nudge_count}...")
                        image_b64 = capture_screen(monitor)
                        messages.append({
                            "role": "user",
                            "content": [
                                {
                                    "type": "text",
                                    "text": (
                                        "You stopped but have not said TASK_COMPLETE. "
                                        "Review the execution plan — are ALL objectives completed and verified? "
                                        "If any remain, continue working. Take a screenshot if needed."
                                    ),
                                },
                                {
                                    "type": "image",
                                    "source": {"type": "base64", "media_type": media_type, "data": image_b64},
                                },
                            ],
                        })

            elif response.stop_reason == "tool_use":
                if tool_results:
                    messages.append({"role": "user", "content": tool_results})
                else:
                    log.warning("tool_use with no results — sending screenshot")
                    image_b64 = capture_screen(monitor)
                    messages.append({
                        "role": "user",
                        "content": [
                            {"type": "text", "text": "Continue with the task:"},
                            {"type": "image", "source": {"type": "base64", "media_type": media_type, "data": image_b64}},
                        ],
                    })

        except KeyboardInterrupt:
            print("\n\n🛑  Interrompido pelo usuário.")
            break
        except anthropic.APIError as e:
            api_retries += 1
            log.error(f"API error (attempt {api_retries}): {e}")
            
            if api_retries >= 5:
                print(f"  ❌  API com falha repetida ({api_retries}x). Abortando.")
                break
            
            # v10: Exponential backoff
            wait_time = min(2 ** api_retries, 30)
            print(f"  ❌  API error: {e} — retry in {wait_time}s...")
            time.sleep(wait_time)
            
            # If context too large, try aggressive compaction
            if "context" in str(e).lower() or "token" in str(e).lower():
                log.info("API error may be context-related — aggressive compaction")
                messages = compact_context(messages, max_images=5)
            
            continue
        except Exception as e:
            log.error(f"Loop error: {e}", exc_info=True)
            print(f"  ❌  Erro: {e}")
            break

    else:
        elapsed = time.time() - start_time
        print(f"\n⚠️   Atingiu o máximo de passos ({MAX_STEPS}) em {elapsed:.0f}s.")
    
    # ── Final stats ────────────────────────────────────────────────────────
    elapsed = time.time() - start_time
    print(f"\n📊  Estatísticas finais:")
    print(f"    Steps: {step} | Tempo: {elapsed:.0f}s")
    print(f"    Tokens: {total_input_tokens:,} input / {total_output_tokens:,} output")
    print(f"    Erros: {len(error_tracker.error_history)}")


# ═══════════════════════════════════════════════════════════════════════════════
# ENTRY POINT
# ═══════════════════════════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(description="Desktop AI Agent v10 — Claude Computer Use")
    parser.add_argument("task", nargs="?", default=None)
    parser.add_argument("--monitor", "-m", type=int, default=DEFAULT_MONITOR)
    parser.add_argument("--list-monitors", "-l", action="store_true")
    parser.add_argument("--no-thinking", action="store_true", help="Disable extended thinking")
    parser.add_argument("--timeout", "-t", type=int, default=TASK_TIMEOUT, help="Task timeout in seconds")
    args = parser.parse_args()

    if args.list_monitors:
        list_monitors()
        sys.exit(0)

    list_monitors()

    task = args.task
    if not task:
        task = input("📝  Digite a tarefa: ").strip()
    if not task:
        task = "Open Notepad and type Hello World"

    run_agent(
        task,
        monitor_index=args.monitor,
        enable_thinking=not args.no_thinking,
        timeout=args.timeout,
    )


if __name__ == "__main__":
    main()