#!/usr/bin/env python3
"""
Otto by Platoona - macOS AI Agent that uses OpenRouter API to control your Mac
"""

import os
import sys
import json
import subprocess
from typing import Optional
from openai import OpenAI

# Configuration - env vars passed by 'otto agent' command
OPENROUTER_KEY = os.getenv("OPENROUTER_KEY")
MODEL = os.getenv("OPENROUTER_MODEL", "openai/gpt-4o")  # Default model
OTTO_PATH = os.path.expanduser("~/.otto/otto")

# OpenRouter client (initialized lazily)
client: Optional[OpenAI] = None

def get_client() -> OpenAI:
    """Get or create the OpenRouter client"""
    global client
    if client is None:
        if not OPENROUTER_KEY:
            print("Error: OpenRouter API key not configured")
            print("\nRun 'otto setup' to configure your API key")
            print("Or get your key from https://openrouter.ai/keys")
            sys.exit(1)
        client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=OPENROUTER_KEY,
        )
    return client

# Tool definitions for all macos-agent commands
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "list_elements",
            "description": "List all clickable UI elements in the frontmost application using Accessibility API. Returns element indices, roles, titles, and positions. Use this as your PRIMARY way to understand what's on screen.",
            "parameters": {
                "type": "object",
                "properties": {
                    "app": {
                        "type": "string",
                        "description": "Optional: Target a specific application by name instead of frontmost app"
                    }
                },
                "required": []
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "click_element",
            "description": "Click a UI element by its index (from list_elements). This is the preferred way to click.",
            "parameters": {
                "type": "object",
                "properties": {
                    "index": {
                        "type": "integer",
                        "description": "Element index from list_elements output"
                    }
                },
                "required": ["index"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "click",
            "description": "Click at specific screen coordinates. Use click_element instead when possible.",
            "parameters": {
                "type": "object",
                "properties": {
                    "x": {
                        "type": "number",
                        "description": "X coordinate"
                    },
                    "y": {
                        "type": "number",
                        "description": "Y coordinate"
                    },
                    "double": {
                        "type": "boolean",
                        "description": "Perform double click"
                    },
                    "right": {
                        "type": "boolean",
                        "description": "Perform right click"
                    }
                },
                "required": ["x", "y"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "type_text",
            "description": "Type text into the currently focused field",
            "parameters": {
                "type": "object",
                "properties": {
                    "text": {
                        "type": "string",
                        "description": "Text to type"
                    }
                },
                "required": ["text"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "press_key",
            "description": "Press a key or key combination (e.g., 'enter', 'cmd+c', 'cmd+space', 'escape')",
            "parameters": {
                "type": "object",
                "properties": {
                    "combo": {
                        "type": "string",
                        "description": "Key or combination like 'enter', 'tab', 'cmd+c', 'cmd+v', 'cmd+space'"
                    }
                },
                "required": ["combo"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "scroll",
            "description": "Scroll at the current mouse position",
            "parameters": {
                "type": "object",
                "properties": {
                    "y": {
                        "type": "integer",
                        "description": "Vertical scroll amount (positive=up, negative=down)"
                    },
                    "x": {
                        "type": "integer",
                        "description": "Horizontal scroll amount (positive=right, negative=left)"
                    },
                    "steps": {
                        "type": "integer",
                        "description": "Number of scroll steps for smooth scrolling"
                    }
                },
                "required": []
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "screenshot",
            "description": "Take a screenshot. Only use as FALLBACK when list_elements doesn't provide enough info.",
            "parameters": {
                "type": "object",
                "properties": {
                    "annotate": {
                        "type": "boolean",
                        "description": "If true, draw bounding boxes with indices on clickable elements"
                    },
                    "output": {
                        "type": "string",
                        "description": "File path to save screenshot (optional)"
                    }
                },
                "required": []
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "run_command",
            "description": "Execute a shell command in the terminal",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "Shell command to execute"
                    },
                    "timeout": {
                        "type": "integer",
                        "description": "Timeout in seconds (default: 30)"
                    }
                },
                "required": ["command"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "wait",
            "description": "Wait for a specified duration before next action",
            "parameters": {
                "type": "object",
                "properties": {
                    "duration": {
                        "type": "string",
                        "description": "Duration like '500ms', '1s', '2.5s'"
                    }
                },
                "required": ["duration"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "list_apps",
            "description": "List all running applications",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "open_app",
            "description": "Open an application by name",
            "parameters": {
                "type": "object",
                "properties": {
                    "app_name": {
                        "type": "string",
                        "description": "Name of the application to open (e.g., 'Calculator', 'Safari', 'Notes')"
                    }
                },
                "required": ["app_name"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "task_complete",
            "description": "Call this when the task is fully completed",
            "parameters": {
                "type": "object",
                "properties": {
                    "summary": {
                        "type": "string",
                        "description": "Brief summary of what was accomplished"
                    }
                },
                "required": ["summary"]
            }
        }
    }
]

SYSTEM_PROMPT = """You are Otto, a macOS automation agent by Platoona. You control the user's computer to accomplish tasks. You MUST be persistent and NEVER give up.

## Tools (in order of preference)
1. list_elements - Get clickable UI elements (TRY THIS FIRST)
2. click_element - Click by element index
3. screenshot (annotate=true) - TAKE THIS when stuck, confused, or list_elements isn't helpful
4. click - Click at x,y coordinates (use after screenshot shows you where)
5. type_text - Type text into focused field
6. press_key - Press keys (enter, escape, cmd+c, tab, etc.)
7. scroll - Scroll up/down
8. wait - Pause for UI to update
9. open_app - Open applications
10. run_command - Run shell commands
11. task_complete - ONLY call when task is FULLY done

## CRITICAL RULES

### Rule 1: ALWAYS take screenshot when stuck
If ANY of these happen, IMMEDIATELY take an annotated screenshot:
- list_elements shows disabled buttons
- You can't find the element you need
- An action didn't work as expected
- You're unsure what's on screen
- A dialog/popup might have appeared
- You've tried the same thing twice without success

### Rule 2: Handle dialogs and popups
Apps often show dialogs like "What's New", "Welcome", "Sign In", etc.
- These block the main UI
- Look for "Continue", "OK", "Skip", "Close", "Not Now" buttons
- Dismiss them before proceeding with the main task

### Rule 3: NEVER give up or ask the user
- Don't say "would you like me to try another approach?"
- Don't stop and ask for guidance
- Take a screenshot, analyze it, and keep trying
- Try different approaches: keyboard shortcuts, clicking coordinates, menu bar

### Rule 4: ALWAYS call task_complete
Every task MUST end with task_complete, either:
- Success: summarize what was accomplished
- Failure: explain what went wrong and what you tried

### Rule 5: Verify your actions
After important actions, check if they worked:
- After clicking: wait 500ms, then list_elements or screenshot
- After typing: verify the text appeared
- After opening app: wait 1-2s for it to fully load

## Strategy
1. Open the app (if needed)
2. Wait 1-2s for app to load
3. list_elements to understand UI
4. If confused or stuck → screenshot with annotate=true
5. Perform action
6. Verify it worked
7. Repeat until done
8. Call task_complete

## Keyboard Shortcuts (use these!)
- cmd+n: New (document, note, window)
- cmd+w: Close window
- cmd+q: Quit app
- escape: Cancel/close dialog
- enter: Confirm/submit
- tab: Next field
"""


def run_otto(args: list[str]) -> str:
    """Run the otto CLI with given arguments"""
    try:
        result = subprocess.run(
            [OTTO_PATH] + args,
            capture_output=True,
            text=True,
            timeout=60
        )
        output = result.stdout
        if result.stderr:
            output += f"\nSTDERR: {result.stderr}"
        if result.returncode != 0:
            output += f"\nExit code: {result.returncode}"
        return output.strip() or "(no output)"
    except subprocess.TimeoutExpired:
        return "Error: Command timed out"
    except Exception as e:
        return f"Error: {str(e)}"


def execute_tool(name: str, args: dict) -> str:
    """Execute a tool and return the result"""

    if name == "list_elements":
        cmd = ["list-elements", "--json"]
        if args.get("app"):
            cmd.extend(["--app", args["app"]])
        return run_otto(cmd)

    elif name == "click_element":
        return run_otto(["click-element", "--index", str(args["index"])])

    elif name == "click":
        cmd = ["click", "--x", str(args["x"]), "--y", str(args["y"])]
        if args.get("double"):
            cmd.append("--double")
        if args.get("right"):
            cmd.append("--right")
        return run_otto(cmd)

    elif name == "type_text":
        return run_otto(["type", args["text"]])

    elif name == "press_key":
        return run_otto(["key", args["combo"]])

    elif name == "scroll":
        cmd = ["scroll"]
        if args.get("y"):
            cmd.extend([f"--y={args['y']}"])
        if args.get("x"):
            cmd.extend([f"--x={args['x']}"])
        if args.get("steps"):
            cmd.extend(["--steps", str(args["steps"])])
        return run_otto(cmd)

    elif name == "screenshot":
        if args.get("annotate"):
            # Smaller image for faster/cheaper API calls
            cmd = ["annotate", "--quality", "40", "--format", "jpeg", "--max-width", "1280", "--max-height", "720"]
        else:
            cmd = ["screenshot", "--format", "jpeg"]

        if args.get("output"):
            # User wants to save to file
            cmd.extend(["--output", args["output"]])
            return run_otto(cmd)
        else:
            # Get base64 directly (no --output means base64 to stdout)
            base64_data = run_otto(cmd)
            if base64_data.startswith("Error"):
                return base64_data
            # Return special marker so we know to include as image
            return f"IMAGE_BASE64:{base64_data}"

    elif name == "run_command":
        cmd = ["run", "--json", "--"]
        cmd.append(args["command"])
        if args.get("timeout"):
            cmd.insert(1, f"--timeout={args['timeout']}")
        return run_otto(cmd)

    elif name == "wait":
        return run_otto(["wait", args["duration"]])

    elif name == "list_apps":
        return run_otto(["list-apps", "--json"])

    elif name == "open_app":
        return run_otto(["run", "--", f"open -a '{args['app_name']}'"])

    elif name == "task_complete":
        return f"TASK_COMPLETE: {args.get('summary', 'Done')}"

    else:
        return f"Unknown tool: {name}"


def run_agent(task: str, max_iterations: int = 50):
    """Run the agent loop for a given task"""

    # Initialize client (will exit if no API key)
    api_client = get_client()

    print(f"\n{'='*60}")
    print(f"Task: {task}")
    print(f"Model: {MODEL}")
    print(f"{'='*60}\n")

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": task}
    ]

    for iteration in range(max_iterations):
        print(f"\n--- Iteration {iteration + 1} ---")

        # Retry logic for transient API errors
        response = None
        for attempt in range(3):
            try:
                response = api_client.chat.completions.create(
                    model=MODEL,
                    messages=messages,
                    tools=TOOLS,
                    tool_choice="auto",
                    extra_headers={
                        "HTTP-Referer": "https://platoona.com/otto",
                        "X-Title": "Otto"
                    }
                )
                break
            except Exception as e:
                error_msg = str(e)
                # Try to extract more details from the error
                if hasattr(e, 'response'):
                    try:
                        error_body = e.response.json() if hasattr(e.response, 'json') else {}
                        if 'error' in error_body:
                            err = error_body['error']
                            error_msg = f"Code {err.get('code', '?')}: {err.get('message', 'Unknown')}"
                            if 'metadata' in err:
                                meta = err['metadata']
                                if 'raw' in meta:
                                    error_msg += f"\n         Raw: {meta['raw'][:300]}"
                                if 'provider_name' in meta:
                                    error_msg += f"\n         Provider: {meta['provider_name']}"
                    except:
                        pass

                print(f"\nAPI Error (attempt {attempt + 1}/3):")
                print(f"  {error_msg}")

                # Log message structure for debugging
                msg_info = []
                for i, m in enumerate(messages):
                    role = m.get('role', '?') if isinstance(m, dict) else getattr(m, 'role', '?')
                    content = m.get('content') if isinstance(m, dict) else getattr(m, 'content', None)
                    if isinstance(content, list):
                        types = [c.get('type', '?') for c in content]
                        msg_info.append(f"{i}:{role}[{'+'.join(types)}]")
                    elif isinstance(content, str):
                        msg_info.append(f"{i}:{role}[text:{len(content)}]")
                    else:
                        msg_info.append(f"{i}:{role}")
                print(f"  Messages: {', '.join(msg_info)}")

                if attempt < 2:
                    import time
                    time.sleep(2)  # Wait before retry
                else:
                    print("\nMax retries reached, stopping.")
                    return

        if response is None:
            break

        message = response.choices[0].message

        # Check if there's text content
        if message.content:
            print(f"\nAgent: {message.content}")

        # Check for tool calls
        if message.tool_calls:
            messages.append(message)

            for tool_call in message.tool_calls:
                func_name = tool_call.function.name
                func_args = json.loads(tool_call.function.arguments) if tool_call.function.arguments else {}

                print(f"\n> Tool: {func_name}")
                print(f"  Args: {json.dumps(func_args, indent=2)}")

                result = execute_tool(func_name, func_args)

                # Truncate long results for display
                if result.startswith("IMAGE_BASE64:"):
                    # Don't print the huge base64 string
                    display_result = "[Screenshot captured - image sent to model]"
                else:
                    display_result = result[:500] + "..." if len(result) > 500 else result
                print(f"  Result: {display_result}")

                # Check if task is complete
                if result.startswith("TASK_COMPLETE:"):
                    print(f"\n{'='*60}")
                    print(result)
                    print(f"{'='*60}\n")
                    return

                # Check if result contains an image (base64)
                if result.startswith("IMAGE_BASE64:"):
                    base64_data = result[len("IMAGE_BASE64:"):]
                    # Tool response must be text, then we add image as user message
                    messages.append({
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": "Screenshot captured successfully. The image is provided below."
                    })
                    # Add image as a user message (OpenAI API requires images in user role)
                    messages.append({
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": "Here is the screenshot you requested. Analyze it and continue with the task:"
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpeg;base64,{base64_data}"
                                }
                            }
                        ]
                    })
                else:
                    messages.append({
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": result
                    })
        else:
            # No tool calls, add message and continue
            messages.append(message)

            # Check if model thinks it's done
            if response.choices[0].finish_reason == "stop":
                print("\nAgent stopped without calling task_complete.")
                break

    print(f"\nMax iterations ({max_iterations}) reached")


def main():
    if len(sys.argv) < 2:
        print("Usage: python agent.py '<task description>'")
        print("\nExamples:")
        print("  python agent.py 'Open Calculator and compute 25 * 4'")
        print("  python agent.py 'Open Notes and create a new note with Hello World'")
        print("  python agent.py 'Take a screenshot and save it to ~/Desktop/screen.png'")
        print("\nEnvironment variables:")
        print("  OPENROUTER_API_KEY - Your OpenRouter API key (required)")
        print("  OPENROUTER_MODEL - Model to use (default: anthropic/claude-sonnet-4)")
        sys.exit(1)

    task = " ".join(sys.argv[1:])
    run_agent(task)


if __name__ == "__main__":
    main()