import anthropic
import base64
import subprocess
from PIL import ImageGrab
import pyautogui
import time
import json
from io import BytesIO

client = anthropic.Anthropic()

def capture_screenshot() -> str:
    """Takes a screenshot and returns base64"""
    screenshot = ImageGrab.grab()
    # Resize to save tokens
    screenshot = screenshot.resize(
        (screenshot.width // 2, screenshot.height // 2)
    )
    buffer = BytesIO()
    screenshot.save(buffer, format="PNG", optimize=True)
    return base64.standard_b64encode(buffer.getvalue()).decode("utf-8")

def execute_computer_action(action: dict) -> str:
    """Executes action on the computer"""
    action_type = action.get("type")

    if action_type == "screenshot":
        return "screenshot_taken"

    elif action_type == "left_click":
        x, y = action["coordinate"]
        pyautogui.click(x * 2, y * 2)
        return f"clicked at ({x}, {y})"

    elif action_type == "double_click":
        x, y = action["coordinate"]
        pyautogui.doubleClick(x * 2, y * 2)
        return f"double-clicked at ({x}, {y})"

    elif action_type == "right_click":
        x, y = action["coordinate"]
        pyautogui.rightClick(x * 2, y * 2)
        return f"right-clicked at ({x}, {y})"

    elif action_type == "type":
        pyautogui.write(action["text"], interval=0.05)
        return f"typed: {action['text'][:50]}..."

    elif action_type == "key":
        pyautogui.hotkey(*action["key"].split("+"))
        return f"pressed: {action['key']}"

    elif action_type == "scroll":
        x, y = action["coordinate"]
        direction = action.get("direction", "down")
        amount = action.get("amount", 3)
        if direction == "down":
            pyautogui.scroll(-amount, x=x * 2, y=y * 2)
        else:
            pyautogui.scroll(amount, x=x * 2, y=y * 2)
        return f"scrolled {direction}"

    elif action_type == "mouse_move":
        x, y = action["coordinate"]
        pyautogui.moveTo(x * 2, y * 2)
        return f"moved to ({x}, {y})"

    return f"unknown action: {action_type}"

class ComputerUseAgent:
    """Agent for GUI automation via Computer Use"""

    COMPUTER_TOOLS = [{
        "type": "computer_20241022",
        "name": "computer",
        "display_width_px": 960,
        "display_height_px": 540,
        "display_number": 1,
    }]

    def __init__(self, max_iterations: int = 50):
        self.max_iterations = max_iterations

    def run_task(self, task: str, context: str = "") -> dict:
        """Executes a task on the computer"""
        system = f"""You control the computer to perform the task.
Use the computer tool to interact with the interface.
Take a screenshot before each action to verify the current state.
Report when the task is complete or if you encounter an insurmountable obstacle.
{"Context: " + context if context else ""}"""

        messages = [{"role": "user", "content": task}]
        iterations = 0
        actions_log = []

        while iterations < self.max_iterations:
            screenshot_b64 = capture_screenshot()

            if messages[-1]["role"] == "user" and isinstance(messages[-1]["content"], str):
                messages[-1]["content"] = [
                    {"type": "text", "text": messages[-1]["content"]},
                    {
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": "image/png",
                            "data": screenshot_b64,
                        },
                    }
                ]

            response = client.beta.messages.create(
                model="claude-opus-4-5",
                max_tokens=4096,
                system=system,
                tools=self.COMPUTER_TOOLS,
                messages=messages,
                betas=["computer-use-2024-10-22"],
            )

            tool_results = []
            task_completed = False

            for block in response.content:
                if hasattr(block, "text"):
                    if any(kw in block.text.lower() for kw in ["task completed", "done", "готово", "completed", "done"]):
                        task_completed = True

                elif block.type == "tool_use" and block.name == "computer":
                    action = block.input
                    result = execute_computer_action(action)
                    actions_log.append({"action": action, "result": result})

                    time.sleep(0.5)

                    new_screenshot = capture_screenshot()
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": [{
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": "image/png",
                                "data": new_screenshot,
                            }
                        }],
                    })

            if task_completed or response.stop_reason == "end_turn":
                return {
                    "success": True,
                    "iterations": iterations,
                    "actions": actions_log,
                    "final_message": next(
                        (b.text for b in response.content if hasattr(b, "text")), ""
                    ),
                }

            messages.append({"role": "assistant", "content": response.content})
            if tool_results:
                messages.append({"role": "user", "content": tool_results})
            else:
                break

            iterations += 1

        return {"success": False, "iterations": iterations, "actions": actions_log}