import base64
import os
import time
import uuid
from typing import Any, cast, Optional, Dict
from PIL import Image
import io

from anthropic import (
    Anthropic,
    AnthropicBedrock,
    AnthropicVertex,
    APIError,
    APIResponseValidationError,
    APIStatusError,
)
from anthropic.types.beta import (
    BetaMessageParam,
    BetaTextBlockParam,
)
from .utils import (
    BATCHED_ACTION_PROMPT,
    BATCHED_TOOL_SCHEMA,
    CLAUDE_47_PROMPT_ADDITIONS,
    PROMPT_CACHING_BETA_FLAG,
    SYSTEM_PROMPT,
    SYSTEM_PROMPT_WINDOWS,
    APIProvider,
    PROVIDER_TO_DEFAULT_MODEL_NAME,
    get_claude_runtime_profile,
    get_model_name,
)
from .utils import _response_to_params, _inject_prompt_caching, _maybe_filter_to_n_most_recent_images
from . import context_policies
from .context_policies import (
    POLICY_KEEP_ALL,
    POLICY_IMAGE_TRIM,
    POLICY_SUMMARIZE,
    POLICY_FOCUS_CHAIN,
    POLICY_FOCUS_CHAIN_V2,
    ALL_POLICIES,
)

import copy
import logging
logger = logging.getLogger("desktopenv.agent")

# MAX_HISTORY = 10
API_RETRY_TIMES = 500  
API_RETRY_INTERVAL = 5

class AnthropicAgent:
    def __init__(self,
                 platform: str = "Ubuntu",
                 model: str = "claude-sonnet-4-5-20250929",
                 provider: APIProvider = APIProvider.BEDROCK,
                 max_tokens: int = 4096,
                 api_key: Optional[str] = None,
                 system_prompt_suffix: str = "",
                 only_n_most_recent_images: Optional[int] = 10,
                 action_space: str = "claude_computer_use",
                 screen_size: tuple[int, int] = (1920, 1080),
                 image_target_size: Optional[tuple[int, int]] = (1280, 720),
                 no_thinking: bool = False,
                 use_isp: bool = False,
                 effort: str = "max",
                 temperature: Optional[float] = None,
                 top_p: Optional[float] = None,
                 context_policy: str = POLICY_KEEP_ALL,
                 context_budget_tokens: int = 25000,
                 keep_n_images: Optional[int] = None,
                 image_chunk_size: Optional[int] = None,
                 *args, **kwargs
                 ):
        self.platform = platform
        self.action_space = action_space
        self.logger = logger
        self.class_name = self.__class__.__name__
        self.model_name = model
        self.provider = provider
        self.max_tokens = max_tokens
        if provider == APIProvider.OPENROUTER:
            self.api_key = api_key
        else:
            self.api_key = api_key or os.environ.get("ANTHROPIC_API_KEY")
        self.system_prompt_suffix = system_prompt_suffix
        self.only_n_most_recent_images = only_n_most_recent_images
        self.messages: list[BetaMessageParam] = []
        self.screen_size = screen_size
        self.image_target_size = image_target_size
        self.no_thinking = no_thinking
        self.use_isp = use_isp
        self.effort = effort
        self.temperature = temperature
        self.top_p = top_p
        self.runtime_profile = get_claude_runtime_profile(model)
        self.max_steps = kwargs.get('max_steps', 15)
        self.current_step = 0
        self.openrouter_session_id: Optional[str] = None

        if context_policy not in ALL_POLICIES:
            raise ValueError(f"context_policy must be one of {ALL_POLICIES}, got {context_policy}")
        self.context_policy = context_policy
        self.context_budget_tokens = context_budget_tokens
        # Per-policy defaults: image_trim keeps more images, trims in small
        # chunks. focus_chain is pure append-only in this experiment
        # (focuschain-ab): notes ride at the end of each user turn and NO
        # screenshots are ever dropped — history is identical to keep_all.
        if context_policy == POLICY_IMAGE_TRIM:
            self.keep_n_images = keep_n_images or 8
            self.image_chunk_size = image_chunk_size or 4
        else:
            self.keep_n_images = keep_n_images
            self.image_chunk_size = image_chunk_size
        if context_policy in (POLICY_KEEP_ALL, POLICY_FOCUS_CHAIN, POLICY_FOCUS_CHAIN_V2):
            # Control arm / pure focus chain: never edit history (the
            # constructor default of 10 would silently trim on long runs).
            self.only_n_most_recent_images = None
            self.keep_n_images = None
        self.latest_task_note: Optional[str] = None
        self.compaction_count = 0
        self._pending_context_usage: Optional[dict] = None
        self._pending_context_event: Optional[str] = None
        # focus_chain_v2 DONE-gate state (per task; also reset in reset())
        self._latest_remaining: list[str] = []
        self._done_gate_fired = False
        self._pending_done_gate_msg: Optional[str] = None
        
        if self.image_target_size:
            self.resize_factor = (
                screen_size[0] / self.image_target_size[0],
                screen_size[1] / self.image_target_size[1],
            )
        else:
            self.resize_factor = (1.0, 1.0)
    
    def _get_sampling_params(self):
        """Get sampling parameters (temperature and/or top_p) - let API validate exclusivity"""
        params = {}
        if self.temperature is not None:
            params['temperature'] = self.temperature
        if self.top_p is not None:
            params['top_p'] = self.top_p
        return params

    def _get_openrouter_api_key(self) -> str:
        # self.api_key silently defaults to ANTHROPIC_API_KEY from the
        # environment; sending an Anthropic key to OpenRouter yields 401
        # "Missing Authentication header". Only trust it if it looks like an
        # OpenRouter key.
        candidate = self.api_key if (self.api_key or "").startswith("sk-or-") else None
        api_key = candidate or os.environ.get("OPENROUTER_API_KEY") or os.environ.get("OPENROUTER_KEY")
        if not api_key:
            raise ValueError("OPENROUTER_API_KEY or OPENROUTER_KEY must be set for OpenRouter")
        return api_key

    def _ensure_openrouter_session_id(self) -> str:
        if not self.openrouter_session_id:
            self.openrouter_session_id = f"osworld-{uuid.uuid4().hex}"
        return self.openrouter_session_id

    def _build_create_kwargs(
        self,
        actual_max_tokens: int,
        system: Any,
        tools: list[dict[str, Any]],
        betas: list[str],
        extra_body: dict[str, Any],
    ) -> dict[str, Any]:
        request_extra_body = dict(extra_body)
        create_kwargs = {
            "max_tokens": actual_max_tokens,
            "messages": self.messages,
            "model": get_model_name(self.provider, self.model_name),
            "system": [system],
            "tools": tools,
            "betas": betas,
            "extra_body": request_extra_body,
            **self._get_sampling_params(),
        }
        if self.provider == APIProvider.OPENROUTER:
            session_id = self._ensure_openrouter_session_id()
            request_extra_body["session_id"] = session_id
            create_kwargs["cache_control"] = {"type": "ephemeral"}
            create_kwargs["extra_headers"] = {"x-session-id": session_id}
        return create_kwargs

    def _extract_usage_metrics(self, response) -> Optional[dict[str, Any]]:
        usage = getattr(response, "usage", None)
        if not usage:
            return None

        raw_usage = usage.model_dump() if hasattr(usage, "model_dump") else {}
        response_extra = getattr(response, "model_extra", None) or {}
        openrouter_metadata = (
            getattr(response, "openrouter_metadata", None)
            or response_extra.get("openrouter_metadata")
        )

        usage_metrics = {
            "input_tokens": getattr(usage, "input_tokens", None),
            "output_tokens": getattr(usage, "output_tokens", None),
            "cache_creation_input_tokens": getattr(usage, "cache_creation_input_tokens", None),
            "cache_read_input_tokens": getattr(usage, "cache_read_input_tokens", None),
            "cache_creation": raw_usage.get("cache_creation"),
            "output_tokens_details": raw_usage.get("output_tokens_details"),
            "service_tier": raw_usage.get("service_tier"),
            "inference_geo": raw_usage.get("inference_geo"),
            "speed": raw_usage.get("speed"),
            "cost": raw_usage.get("cost"),
            "cost_details": raw_usage.get("cost_details"),
            "is_byok": raw_usage.get("is_byok"),
            "provider": getattr(response, "provider", None) or response_extra.get("provider"),
            "openrouter_metadata": openrouter_metadata,
            "openrouter_session_id": self.openrouter_session_id,
            "actual_model": getattr(response, "model", None),
        }
        usage_metrics["context_policy"] = self.context_policy
        if self._pending_context_event:
            usage_metrics["context_event"] = self._pending_context_event
            self._pending_context_event = None
        if self._pending_context_usage:
            usage_metrics["context_mgmt_usage"] = self._pending_context_usage
            self._pending_context_usage = None
        return {key: value for key, value in usage_metrics.items() if value is not None}

    def _focus_chain_v2_tool_schema(self) -> dict:
        """BATCHED_TOOL_SCHEMA plus the required structured v2 ledger."""
        schema = copy.deepcopy(BATCHED_TOOL_SCHEMA)
        schema["input_schema"]["properties"]["task_notes"] = copy.deepcopy(
            context_policies.TASK_NOTES_V2_SCHEMA_PROPERTY
        )
        required = schema["input_schema"].setdefault("required", [])
        if "task_notes" not in required:
            required.append("task_notes")
        return schema

    def _focus_chain_tool_schema(self) -> dict:
        """BATCHED_TOOL_SCHEMA plus a required running task_notes field (P3)."""
        schema = copy.deepcopy(BATCHED_TOOL_SCHEMA)
        schema["input_schema"]["properties"]["task_notes"] = copy.deepcopy(
            context_policies.TASK_NOTES_SCHEMA_PROPERTY
        )
        required = schema["input_schema"].setdefault("required", [])
        if "task_notes" not in required:
            required.append("task_notes")
        return schema

    def add_tool_result(self, tool_call_id: str, result: str, screenshot: bytes = None):
        """Add tool result to message history"""
        tool_result_content = [
            {
                "type": "tool_result",
                "tool_use_id": tool_call_id,
                "content": [{"type": "text", "text": result}]
            }
        ]
        
        # Add screenshot if provided
        if screenshot is not None:
            screenshot_base64 = base64.b64encode(screenshot).decode('utf-8')
            tool_result_content[0]["content"].append({
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/png", 
                    "data": screenshot_base64
                }
            })
        
        self.messages.append({
            "role": "user",
            "content": tool_result_content
        })
    
    def _extract_raw_response_string(self, response) -> str:
        """Extract and concatenate raw response content into a single string."""
        raw_response_str = ""
        if response.content:
            for block in response.content:
                if hasattr(block, 'text') and block.text:
                    raw_response_str += f"[TEXT] {block.text}\n"
                elif hasattr(block, 'thinking') and block.thinking:
                    raw_response_str += f"[THINKING] {block.thinking}\n"
                elif hasattr(block, 'name') and hasattr(block, 'input'):
                    raw_response_str += f"[TOOL_USE] {block.name}: {block.input}\n"
                else:
                    raw_response_str += f"[OTHER] {str(block)}\n"
        return raw_response_str.strip()

    def parse_actions_from_tool_call(self, tool_call: Dict) -> str:
        result = ""
        function_args = (
            tool_call["input"]
        )

        batched_actions = function_args.get("actions")
        if batched_actions is not None:
            if not isinstance(batched_actions, list):
                raise ValueError(f"{batched_actions} must be a list")
            for batched_action in batched_actions:
                if not isinstance(batched_action, dict):
                    raise ValueError(f"{batched_action} must be a dict")
                result += self.parse_actions_from_tool_call({"input": batched_action})
            return result
        
        action = function_args.get("action")
        if not action:
            action = tool_call.function.name
        action_conversion = {
            "left click": "click",
            "right click": "right_click"
        }
        action = action_conversion.get(action, action)
        
        text = function_args.get("text")
        coordinate = function_args.get("coordinate")
        start_coordinate = function_args.get("start_coordinate")
        scroll_direction = function_args.get("scroll_direction")
        scroll_amount = function_args.get("scroll_amount")
        duration = function_args.get("duration")
        repeat = int(function_args.get("repeat") or 1)
        
        # resize coordinates if resize_factor is set
        if coordinate and self.resize_factor:
            coordinate = (
                int(coordinate[0] * self.resize_factor[0]),
                int(coordinate[1] * self.resize_factor[1])
            )
        if start_coordinate and self.resize_factor:
            start_coordinate = (
                int(start_coordinate[0] * self.resize_factor[0]),
                int(start_coordinate[1] * self.resize_factor[1])
            )
        
        if action == "left_mouse_down":
            result += "pyautogui.mouseDown()\n"
        elif action == "left_mouse_up":
            result += "pyautogui.mouseUp()\n"
        
        elif action == "hold_key":
            if not isinstance(text, str):
                raise ValueError(f"{text} must be a string")
            
            keys = text.split('+')
            for key in keys:
                key = key.strip().lower()
                result += f"pyautogui.keyDown('{key}')\n"
            expected_outcome = f"Keys {text} held down."

        # Handle mouse move and drag actions
        elif action in ("mouse_move", "left_click_drag"):
            if coordinate is None:
                raise ValueError(f"coordinate is required for {action}")
            if text is not None:
                raise ValueError(f"text is not accepted for {action}")
            if not isinstance(coordinate, (list, tuple)) or len(coordinate) != 2:
                raise ValueError(f"{coordinate} must be a tuple of length 2")
            if not all(isinstance(i, int) for i in coordinate):
                raise ValueError(f"{coordinate} must be a tuple of ints")
            
            x, y = coordinate[0], coordinate[1]
            if action == "mouse_move":
                result += (
                    f"pyautogui.moveTo({x}, {y}, duration={duration or 0.5})\n"
                )
                expected_outcome = f"Mouse moved to ({x},{y})."
            elif action == "left_click_drag":
                # If start_coordinate is provided, validate and move to start before dragging
                if start_coordinate:
                    if not isinstance(start_coordinate, (list, tuple)) or len(start_coordinate) != 2:
                        raise ValueError(f"{start_coordinate} must be a tuple of length 2")
                    if not all(isinstance(i, int) for i in start_coordinate):
                        raise ValueError(f"{start_coordinate} must be a tuple of ints")
                    start_x, start_y = start_coordinate[0], start_coordinate[1]
                    result += (
                        f"pyautogui.moveTo({start_x}, {start_y}, duration={duration or 0.5})\n"
                    )
                result += (
                    f"pyautogui.dragTo({x}, {y}, duration={duration or 0.5})\n"
                )
                expected_outcome = f"Cursor dragged to ({x},{y})."

        # Handle keyboard actions
        elif action in ("key", "type"):
            if text is None:
                raise ValueError(f"text is required for {action}")
            if coordinate is not None:
                raise ValueError(f"coordinate is not accepted for {action}")
            if not isinstance(text, str):
                raise ValueError(f"{text} must be a string")

            if action == "key":
                key_conversion = {
                    "page_down": "pagedown",
                    "page_up": "pageup",
                    "super_l": "win",
                    "super": "command",
                    "escape": "esc"
                }
                keys = text.split('+')
                for _ in range(repeat):
                    for key in keys:
                        key = key.strip().lower()
                        key = key_conversion.get(key, key)
                        result += (f"pyautogui.keyDown('{key}')\n")
                    for key in reversed(keys):
                        key = key.strip().lower()
                        key = key_conversion.get(key, key)
                        result += (f"pyautogui.keyUp('{key}')\n")
                expected_outcome = f"Key {text} pressed {repeat} time(s)."
            elif action == "type":
                escaped_text = repr(text)
                paste_script = [
                    "try:",
                    "    import pyperclip",
                    f"    pyperclip.copy({escaped_text})",
                    "    pyautogui.hotkey('ctrl', 'v')",
                    "except Exception:",
                ]
                for char in text:
                    if char == '\n':
                        paste_script.append("    pyautogui.press('enter')")
                    elif char == "'":
                        paste_script.append('    pyautogui.press("\'")')
                    elif char == '\\':
                        paste_script.append("    pyautogui.press('\\\\')")
                    elif char == '"':
                        paste_script.append("    pyautogui.press('\"')")
                    else:
                        paste_script.append(f"    pyautogui.press({repr(char)})")
                result += f"exec({repr(chr(10).join(paste_script))})\n"
                expected_outcome = f"Text {text} written."

        # Handle scroll actions
        elif action == "scroll":
            if text is not None:
                result += (f"pyautogui.keyDown('{text.lower()}')\n")
            if coordinate is None:
                if scroll_direction in ("up", "down"):
                    result += (
                        f"pyautogui.scroll({scroll_amount if scroll_direction == 'up' else -scroll_amount})\n"
                    )
                elif scroll_direction in ("left", "right"):
                    result += (
                        f"pyautogui.hscroll({scroll_amount if scroll_direction == 'right' else -scroll_amount})\n"
                    )
            else:
                if scroll_direction in ("up", "down"):
                    x, y = coordinate[0], coordinate[1]
                    result += (
                        f"pyautogui.scroll({scroll_amount if scroll_direction == 'up' else -scroll_amount}, {x}, {y})\n"
                    )
                elif scroll_direction in ("left", "right"):
                    x, y = coordinate[0], coordinate[1]
                    result += (
                        f"pyautogui.hscroll({scroll_amount if scroll_direction == 'right' else -scroll_amount}, {x}, {y})\n"
                    )
            if text is not None:
                result += (f"pyautogui.keyUp('{text.lower()}')\n")
            expected_outcome = "Scroll action finished"

        # Handle click actions
        elif action in ("left_click", "right_click", "double_click", "middle_click", "left_press", "triple_click"):
            # Handle modifier keys during click if specified
            if text:
                keys = text.split('+')
                for key in keys:
                    key = key.strip().lower()
                    result += f"pyautogui.keyDown('{key}')\n"
            if coordinate is not None:
                x, y = coordinate
                if action == "left_click":
                    result += (f"pyautogui.click({x}, {y})\n")
                elif action == "right_click":
                    result += (f"pyautogui.rightClick({x}, {y})\n")
                elif action == "double_click":
                    result += (f"pyautogui.doubleClick({x}, {y})\n")
                elif action == "middle_click":
                    result += (f"pyautogui.middleClick({x}, {y})\n")
                elif action == "left_press":
                    result += (f"pyautogui.mouseDown({x}, {y})\n")
                    result += ("time.sleep(1)\n")
                    result += (f"pyautogui.mouseUp({x}, {y})\n")
                elif action == "triple_click":
                    result += (f"pyautogui.tripleClick({x}, {y})\n")

            else:
                if action == "left_click":
                    result += ("pyautogui.click()\n")
                elif action == "right_click":
                    result += ("pyautogui.rightClick()\n")
                elif action == "double_click":
                    result += ("pyautogui.doubleClick()\n")
                elif action == "middle_click":
                    result += ("pyautogui.middleClick()\n")
                elif action == "left_press":
                    result += ("pyautogui.mouseDown()\n")
                    result += ("time.sleep(1)\n")
                    result += ("pyautogui.mouseUp()\n")
                elif action == "triple_click":
                    result += ("pyautogui.tripleClick()\n")
            # Release modifier keys after click
            if text:
                keys = text.split('+')
                for key in reversed(keys):
                    key = key.strip().lower()
                    result += f"pyautogui.keyUp('{key}')\n"
            expected_outcome = "Click action finished"
            
        elif action == "wait":
            result += f"pyautogui.sleep({duration or 0.5})\n"
            expected_outcome = f"Wait for {duration or 0.5} seconds"
        elif action == "fail":
            result += "FAIL"
            expected_outcome = "Finished"
        elif action == "done":
            result += "DONE"
            expected_outcome = "Finished"
        elif action == "call_user":
            result += "CALL_USER"
            expected_outcome = "Call user"
        elif action == "screenshot":
            result += "pyautogui.sleep(0.1)\n"
            expected_outcome = "Screenshot taken"
        else:
            raise ValueError(f"Invalid action: {action}")
        
        return result
            
    def predict(self, task_instruction: str, obs: Dict = None, system: Any = None):
        self.current_step += 1

        system_text = f"{SYSTEM_PROMPT_WINDOWS if self.platform == 'Windows' else SYSTEM_PROMPT}{' ' + self.system_prompt_suffix if self.system_prompt_suffix else ''}"
        if self.runtime_profile.use_claude_47_prompt:
            system_text += f"\n\n{CLAUDE_47_PROMPT_ADDITIONS}"
            system_text += f"\n\n<STEP_LIMIT>You have a maximum of {self.max_steps} steps to complete this task. Use your steps wisely.</STEP_LIMIT>"
            system_text += f"\n\n{BATCHED_ACTION_PROMPT}"
        elif self.provider == APIProvider.OPENROUTER:
            system_text += f"\n\n{CLAUDE_47_PROMPT_ADDITIONS}"
            system_text += f"\n\n<STEP_LIMIT>You have a maximum of {self.max_steps} steps to complete this task. Use your steps wisely.</STEP_LIMIT>"
            system_text += f"\n\n{BATCHED_ACTION_PROMPT}"

        if self.context_policy == POLICY_SUMMARIZE:
            system_text += f"\n\n{context_policies.SUMMARIZE_PROMPT_ADDITION}"
        elif self.context_policy == POLICY_FOCUS_CHAIN:
            system_text += f"\n\n{context_policies.FOCUS_CHAIN_PROMPT_ADDITION}"
        elif self.context_policy == POLICY_FOCUS_CHAIN_V2:
            system_text += f"\n\n{context_policies.FOCUS_CHAIN_V2_PROMPT_ADDITION}"

        system = BetaTextBlockParam(
            type="text",
            text=system_text,
        )
        
        image_metrics = None

        # Resize screenshot to the configured computer-use display size.
        if obs and "screenshot" in obs:
            # Convert bytes to PIL Image
            screenshot_bytes = obs["screenshot"]
            screenshot_image = Image.open(io.BytesIO(screenshot_bytes))
            
            # Store original unresized screenshot for zoom processing
            obs["screenshot_original"] = screenshot_bytes
            
            original_width, original_height = screenshot_image.size
            if self.image_target_size:
                new_width, new_height = self.image_target_size

                resized_image = screenshot_image.resize((new_width, new_height), Image.Resampling.LANCZOS)

                output_buffer = io.BytesIO()
                resized_image.save(output_buffer, format='PNG')
                obs["screenshot"] = output_buffer.getvalue()
            else:
                new_width, new_height = original_width, original_height

            image_metrics = {
                "original_width": original_width,
                "original_height": original_height,
                "sent_width": new_width,
                "sent_height": new_height,
                "sent_bytes": len(obs["screenshot"]),
            }
            

        if not self.messages:
            
            init_screenshot = obs
            init_screenshot_base64 = base64.b64encode(init_screenshot["screenshot"]).decode('utf-8')
            self.messages.append({
                "role": "user",
                "content": [
                    {
                    "type": "image",
                    "source": {
                            "type": "base64",
                            "media_type": "image/png",
                            "data": init_screenshot_base64,
                        },
                    },
                    {"type": "text", "text": task_instruction},
                ]
            })
            
        # v2 DONE-gate delivery: the rejected completion was a plain-text
        # assistant message (no tool_use), so no tool_result turn will be
        # created below. Append a plain user turn: fresh screenshot + the
        # rejection + the ledger echo. Append-only.
        if (
            self.context_policy == POLICY_FOCUS_CHAIN_V2
            and self._pending_done_gate_msg
            and self.messages
            and self.messages[-1]["role"] == "assistant"
            and obs and obs.get("screenshot")
        ):
            gate_content = [
                {
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": "image/png",
                        "data": base64.b64encode(obs["screenshot"]).decode("utf-8"),
                    },
                },
                {"type": "text", "text": self._pending_done_gate_msg},
            ]
            if self.latest_task_note:
                gate_content.append({
                    "type": "text",
                    "text": f"[TASK_STATE_NOTES from your previous step]\n{self.latest_task_note}",
                })
            self.messages.append({"role": "user", "content": gate_content})
            self._pending_done_gate_msg = None

        # Add tool_result for ALL tool_use blocks in the last message
        if self.messages:
            last_message_content = self.messages[-1]["content"]
            tool_use_blocks = [block for block in last_message_content if block.get("type") == "tool_use"]
            
            for i, tool_block in enumerate(tool_use_blocks):
                tool_input = tool_block.get("input", {})
                action = tool_input.get("action")
                is_last_tool = i == len(tool_use_blocks) - 1
                
                include_screenshot = None
                
                if obs:
                    if action == "screenshot":
                        # Screenshot action always gets regular screenshot
                        include_screenshot = obs.get("screenshot")
                    elif is_last_tool:
                        # Auto-screenshot: last tool gets regular screenshot (unless it's zoom, handled above)
                        include_screenshot = obs.get("screenshot")
                
                self.add_tool_result(
                    tool_block["id"],
                    f"Success",
                    screenshot=include_screenshot
                )

            if tool_use_blocks and self.runtime_profile.use_claude_47_prompt and self.messages and self.messages[-1]["role"] == "user":
                self.messages[-1]["content"].append({
                    "type": "text",
                    "text": f"[Current step: {self.current_step}/{self.max_steps}]",
                })

            # P3 focus_chain: append the model's own latest task-state note at
            # the END of the new user turn. Append-only — earlier history text
            # is never rewritten, so the cached prefix stays valid.
            if (
                self.context_policy in (POLICY_FOCUS_CHAIN, POLICY_FOCUS_CHAIN_V2)
                and tool_use_blocks
                and self.latest_task_note
                and self.messages
                and self.messages[-1]["role"] == "user"
            ):
                self.messages[-1]["content"].append({
                    "type": "text",
                    "text": f"[TASK_STATE_NOTES from your previous step]\n{self.latest_task_note}",
                })

        enable_prompt_caching = False
        betas = [self.runtime_profile.computer_use_beta_flag]

        # Claude 4.5/4.6 keep the legacy ISP path. Claude 4.7 uses adaptive effort.
        if self.use_isp and self.runtime_profile.thinking_mode == "legacy":
            betas.append("interleaved-thinking-2025-05-14")
            logger.info(f"Added interleaved thinking beta. Betas: {betas}")
            
        image_truncation_threshold = 10
        if self.provider == APIProvider.ANTHROPIC:
            client = Anthropic(api_key=self.api_key, max_retries=4).with_options(
                default_headers={"anthropic-beta": self.runtime_profile.computer_use_beta_flag}
            )
            enable_prompt_caching = True
        elif self.provider == APIProvider.OPENROUTER:
            openrouter_key = self._get_openrouter_api_key()
            client = Anthropic(
                api_key=openrouter_key,
                base_url=os.environ.get("OPENROUTER_BASE_URL", "https://openrouter.ai/api"),
                max_retries=4,
                default_headers={
                    "Authorization": f"Bearer {openrouter_key}",
                    "X-OpenRouter-Metadata": "enabled",
                    "X-Title": "OSWorld Computer Use Eval",
                },
            )
            enable_prompt_caching = True
        elif self.provider == APIProvider.VERTEX:
            client = AnthropicVertex()
        elif self.provider == APIProvider.BEDROCK:
            client = AnthropicBedrock(
                # Authenticate by either providing the keys below or use the default AWS credential providers, such as
                # using ~/.aws/credentials or the "AWS_SECRET_ACCESS_KEY" and "AWS_ACCESS_KEY_ID" environment variables.
                aws_access_key=os.getenv('AWS_ACCESS_KEY_ID'),
                aws_secret_key=os.getenv('AWS_SECRET_ACCESS_KEY'),
                # aws_region changes the aws region to which the request is made. By default, we read AWS_REGION,
                # and if that's not present, we default to us-east-1. Note that we do not read ~/.aws/config for the region.
                aws_region=os.getenv('AWS_DEFAULT_REGION'),
            )

        if enable_prompt_caching:
            betas.append(PROMPT_CACHING_BETA_FLAG)
            image_truncation_threshold = 20
            if self.provider == APIProvider.OPENROUTER:
                self._ensure_openrouter_session_id()
            else:
                _inject_prompt_caching(self.messages)
                system["cache_control"] = {"type": "ephemeral"}

        # --- Context-management policy dispatch (cache-aware-context experiment) ---
        self._pending_context_event = None
        if self.context_policy == POLICY_IMAGE_TRIM:
            images_before = context_policies.count_images(self.messages)
            _maybe_filter_to_n_most_recent_images(
                self.messages,
                self.keep_n_images,
                min_removal_threshold=self.image_chunk_size,
            )
            images_after = context_policies.count_images(self.messages)
            if images_after < images_before:
                self._pending_context_event = (
                    f"image_trim_fired:removed={images_before - images_after}"
                )
                logger.info(
                    f"[{self.context_policy}] trimmed {images_before - images_after} images "
                    f"({images_before} -> {images_after})"
                )
        elif self.context_policy == POLICY_SUMMARIZE:
            estimated = context_policies.estimate_tokens(self.messages, system_text)
            threshold = int(0.6 * self.context_budget_tokens)
            if estimated > threshold:
                logger.info(
                    f"[summarize] estimated {estimated} tokens > threshold {threshold}, compacting"
                )
                new_messages, summarizer_usage, did_compact = context_policies.compact_messages(
                    self.messages,
                    client,
                    get_model_name(self.provider, self.model_name),
                    task_instruction,
                )
                if did_compact:
                    self.messages = new_messages
                    self.compaction_count += 1
                    self._pending_context_usage = summarizer_usage
                    self._pending_context_event = (
                        f"compaction_fired:count={self.compaction_count},est_tokens={estimated}"
                    )
        elif self.only_n_most_recent_images:
            # Legacy path (keep_all leaves history untouched; only reachable if
            # explicitly configured outside the experiment policies).
            _maybe_filter_to_n_most_recent_images(
                self.messages,
                self.only_n_most_recent_images,
                min_removal_threshold=image_truncation_threshold,
            )

        # OpenRouter currently rejects Anthropic's built-in computer-use tool
        # shorthand, so use the custom schema that the local parser already
        # understands. Direct Anthropic/Bedrock/Vertex keep native computer use.
        if self.provider == APIProvider.OPENROUTER:
            if self.context_policy == POLICY_FOCUS_CHAIN:
                tools = [self._focus_chain_tool_schema()]
            elif self.context_policy == POLICY_FOCUS_CHAIN_V2:
                tools = [self._focus_chain_v2_tool_schema()]
            else:
                tools = [BATCHED_TOOL_SCHEMA]
        else:
            tool_config = {
                'name': 'computer',
                'type': self.runtime_profile.computer_use_type,
                'display_width_px': self.image_target_size[0] if self.image_target_size else self.screen_size[0],
                'display_height_px': self.image_target_size[1] if self.image_target_size else self.screen_size[1],
                'display_number': 1
            }
            tools = [tool_config]
        
        # Configure thinking mode by Claude version.
        if self.runtime_profile.thinking_mode == "adaptive":
            extra_body = {
                "thinking": {"type": "adaptive"},
                "output_config": {"effort": self.effort},
            }
            actual_max_tokens = max(self.max_tokens, self.runtime_profile.default_max_tokens)
            logger.info(f"Runtime profile: {self.runtime_profile.label}; thinking mode: ADAPTIVE; effort: {self.effort}")
        else:
            if self.no_thinking:
                # Disable thinking mode - omit the thinking parameter.
                extra_body = {}
                actual_max_tokens = self.max_tokens
                logger.info(f"Runtime profile: {self.runtime_profile.label}; thinking mode: DISABLED")
            else:
                # Legacy 4.5/4.6 behavior: fixed extended-thinking budget.
                budget_tokens = 2048

                # For regular thinking: max_tokens > budget_tokens (API requirement)
                # For ISP: budget_tokens can exceed max_tokens (represents total across all thinking blocks)
                if self.max_tokens <= budget_tokens:
                    required_max_tokens = budget_tokens + 500  # Give some headroom
                    logger.warning(f"Regular thinking requires max_tokens > budget_tokens. Increasing max_tokens from {self.max_tokens} to {required_max_tokens}")
                    actual_max_tokens = required_max_tokens
                else:
                    actual_max_tokens = self.max_tokens

                extra_body = {
                    "thinking": {"type": "enabled", "budget_tokens": budget_tokens}
                }
                if self.use_isp:
                    logger.info(f"Runtime profile: {self.runtime_profile.label}; thinking mode: INTERLEAVED SCRATCHPAD (ISP)")
                else:
                    logger.info(f"Runtime profile: {self.runtime_profile.label}; thinking mode: REGULAR SCRATCHPAD")

        try:
            response = None
            
            for attempt in range(API_RETRY_TIMES):
                try:
                    response = client.beta.messages.create(
                        **self._build_create_kwargs(
                            actual_max_tokens=actual_max_tokens,
                            system=system,
                            tools=tools,
                            betas=betas,
                            extra_body=extra_body,
                        )
                    )
                    
                    logger.info(f"Response: {response}")
                    break  
                except (APIError, APIStatusError, APIResponseValidationError) as e:
                    error_msg = str(e)
                    logger.warning(f"Anthropic API error (attempt {attempt+1}/{API_RETRY_TIMES}): {error_msg}")
                    
                    # Detect payload size errors (413 request_too_large, 25MB limit, etc.)
                    is_size_error = (
                        "25000000" in error_msg or 
                        "Member must have length less than or equal to" in error_msg or
                        "request_too_large" in error_msg or
                        "maximum size" in error_msg.lower() or
                        "413" in error_msg
                    )
                    
                    if is_size_error:
                        logger.warning("Detected request size error (413/25MB limit), automatically reducing image count")
                        current_image_count = self.only_n_most_recent_images or context_policies.count_images(self.messages)
                        new_image_count = max(1, current_image_count // 2)  # Keep at least 1 image
                        self.only_n_most_recent_images = new_image_count
                        
                        _maybe_filter_to_n_most_recent_images(
                            self.messages,
                            new_image_count,
                            min_removal_threshold=image_truncation_threshold,
                        )
                        logger.info(f"Image count reduced from {current_image_count} to {new_image_count}")
                    
                    if attempt < API_RETRY_TIMES - 1:
                        time.sleep(API_RETRY_INTERVAL)
                    else:
                        raise  # All attempts failed, raise exception to enter existing except logic

        except (APIError, APIStatusError, APIResponseValidationError) as e:
            logger.exception(f"Anthropic API error: {str(e)}")
            if self.provider != APIProvider.ANTHROPIC:
                return None, None
            try:
                logger.warning("Retrying with backup API key...")

                backup_client = Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY_BACKUP"), max_retries=4).with_options(
                    default_headers={"anthropic-beta": self.runtime_profile.computer_use_beta_flag}
                )
                response = backup_client.beta.messages.create(
                    **self._build_create_kwargs(
                        actual_max_tokens=actual_max_tokens,
                        system=system,
                        tools=tools,
                        betas=betas,
                        extra_body=extra_body,
                    )
                )
                
                logger.info("Successfully used backup API key")
            except Exception as backup_e:
                backup_error_msg = str(backup_e)
                logger.exception(f"Backup API call also failed: {backup_error_msg}")
                
                # Check if backup API also has request size error (413, 25MB limit, etc.)
                is_size_error = (
                    "25000000" in backup_error_msg or 
                    "Member must have length less than or equal to" in backup_error_msg or
                    "request_too_large" in backup_error_msg or
                    "maximum size" in backup_error_msg.lower() or
                    "413" in backup_error_msg
                )
                
                if is_size_error:
                    logger.warning("Backup API also encountered request size error (413/25MB limit), further reducing image count")
                    # Reduce image count by half again
                    current_image_count = self.only_n_most_recent_images or context_policies.count_images(self.messages)
                    new_image_count = max(1, current_image_count // 2)  # Keep at least 1 image
                    self.only_n_most_recent_images = new_image_count
                    
                    # Reapply image filtering
                    _maybe_filter_to_n_most_recent_images(
                        self.messages,
                        new_image_count,
                        min_removal_threshold=image_truncation_threshold,
                    )
                    logger.info(f"Backup API image count reduced from {current_image_count} to {new_image_count}")
                
                return None, None

        except Exception as e:
            logger.exception(f"Error in Anthropic API: {str(e)}")
            return None, None

        if response is None:
            logger.error("Response is None after API call - this should not happen")
            return None, None

        response_params = _response_to_params(response)
        logger.info(f"Received response params: {response_params}")

        usage_metrics = self._extract_usage_metrics(response)

        # Convert raw response to concatenated string for trajectory logging
        raw_response_str = self._extract_raw_response_string(response)

        # Store response in message history
        self.messages.append({
            "role": "assistant",
            "content": response_params
        })

        # P3 focus_chain: capture the model's running task-state note.
        if self.context_policy == POLICY_FOCUS_CHAIN:
            for content_block in response_params:
                if content_block.get("type") == "tool_use":
                    note = (content_block.get("input") or {}).get("task_notes")
                    if isinstance(note, str) and note.strip():
                        self.latest_task_note = note.strip()
        elif self.context_policy == POLICY_FOCUS_CHAIN_V2:
            for content_block in response_params:
                if content_block.get("type") == "tool_use":
                    note = (content_block.get("input") or {}).get("task_notes")
                    if isinstance(note, dict) and note:
                        self.latest_task_note = context_policies.render_task_notes_v2(note)
                        self._latest_remaining = [
                            str(r) for r in (note.get("remaining") or []) if str(r).strip()
                        ]

        max_parse_retry = 3
        for parse_retry in range(max_parse_retry):
            actions: list[Any] = []
            reasonings: list[str] = []
            try:
                for content_block in response_params:
                    if content_block["type"] == "tool_use":
                        actions.append({
                            "name": content_block["name"],
                            "input": cast(dict[str, Any], content_block["input"]),
                            "id": content_block["id"],
                            "action_type": content_block.get("type"),
                            "command": self.parse_actions_from_tool_call(content_block),
                            "raw_response": raw_response_str,
                            "usage": usage_metrics,
                            "image_metrics": image_metrics,
                        })
                    elif content_block["type"] == "text":
                        reasonings.append(content_block["text"])
                if isinstance(reasonings, list) and len(reasonings) > 0:
                    reasonings = reasonings[0]
                else:
                    reasonings = ""
                
                # Check if the model indicated the task is infeasible
                if raw_response_str and "[INFEASIBLE]" in raw_response_str:
                    logger.info("Detected [INFEASIBLE] pattern in response, triggering FAIL action")
                    # Override actions with FAIL
                    actions = [{
                        "action_type": "FAIL",
                        "raw_response": raw_response_str,
                        "usage": usage_metrics,
                        "image_metrics": image_metrics,
                    }]
                
                logger.info(f"Received actions: {actions}")
                logger.info(f"Received reasonings: {reasonings}")
                if len(actions) == 0:
                    # A response with no tool call is this stack's DONE signal.
                    # v2 DONE-gate: reject it once if the model's own ledger
                    # still lists unverified items; force a fresh observation
                    # instead. History is untouched (append-only); the model
                    # sees its attempt + the rejection next turn.
                    if (
                        self.context_policy == POLICY_FOCUS_CHAIN_V2
                        and not self._done_gate_fired
                        and self._latest_remaining
                    ):
                        self._done_gate_fired = True
                        remaining_txt = "; ".join(self._latest_remaining[:6])
                        self._pending_done_gate_msg = (
                            "[DONE_GATE] Completion NOT accepted: your task ledger "
                            f"still lists unverified items: {remaining_txt}. "
                            "A fresh screenshot follows. Verify each item on screen "
                            "(move it to verified_done) or finish it. Your next "
                            "done declaration will be accepted."
                        )
                        self._pending_context_event = (
                            f"done_gate_fired:remaining={len(self._latest_remaining)}"
                        )
                        logger.info(
                            f"[focus_chain_v2] DONE-gate rejected completion; "
                            f"remaining={self._latest_remaining}"
                        )
                        actions = [{
                            "action_type": "tool_use",
                            "name": "computer",
                            "id": "done_gate_forced_screenshot",
                            "input": {"actions": [{"action": "screenshot"}]},
                            "command": "pyautogui.sleep(0.1)\n",
                            "raw_response": raw_response_str,
                            "usage": usage_metrics,
                            "image_metrics": image_metrics,
                        }]
                    else:
                        actions = [{
                            "action_type": "DONE",
                            "raw_response": raw_response_str,
                            "usage": usage_metrics,
                            "image_metrics": image_metrics,
                        }]
                return reasonings, actions
            except Exception as e:
                logger.warning(f"parse_actions_from_tool_call parsing failed (attempt {parse_retry+1}/3), will retry API request: {e}")
                # Remove the recently appended assistant message to avoid polluting history
                self.messages.pop()
                # Retry API request
                response = None
                for attempt in range(API_RETRY_TIMES):
                    try:
                        response = client.beta.messages.create(
                            **self._build_create_kwargs(
                                actual_max_tokens=actual_max_tokens,
                                system=system,
                                tools=tools,
                                betas=betas,
                                extra_body=extra_body,
                            )
                        )
                        
                        logger.info(f"Response: {response}")
                        break  # Success, exit retry loop
                    except (APIError, APIStatusError, APIResponseValidationError) as e2:
                        error_msg = str(e2)
                        logger.warning(f"Anthropic API error (attempt {attempt+1}/{API_RETRY_TIMES}): {error_msg}")
                        if attempt < API_RETRY_TIMES - 1:
                            time.sleep(API_RETRY_INTERVAL)
                        else:
                            raise
                response_params = _response_to_params(response)
                logger.info(f"Received response params: {response_params}")
                
                # Update raw response string for retry case (will be used in next loop iteration)
                raw_response_str = self._extract_raw_response_string(response)
                usage_metrics = self._extract_usage_metrics(response)
                
                self.messages.append({
                    "role": "assistant",
                    "content": response_params
                })
                if parse_retry == max_parse_retry - 1:
                    logger.error(f"parse_actions_from_tool_call parsing failed 3 times consecutively, terminating: {e}")
                    actions = [{
                        "action_type": "FAIL",
                        "raw_response": f"Failed to parse actions from tool call after {max_parse_retry} attempts: {e}",
                        "usage": usage_metrics,
                        "image_metrics": image_metrics,
                    }]
                    return reasonings, actions
    def reset(self, _logger = None, *args, **kwargs):
        """
        Reset the agent's state.
        """
        global logger
        if _logger:
            logger = _logger
        else:
            logger = logging.getLogger("desktopenv.agent")
        self.messages = []
        self.current_step = 0
        self.openrouter_session_id = None
        self.latest_task_note = None
        self.compaction_count = 0
        self._pending_context_usage = None
        self._pending_context_event = None
        self._latest_remaining = []
        self._done_gate_fired = False
        self._pending_done_gate_msg = None
        logger.info(f"{self.class_name} reset.")
