"""
Utility functions for the Anthropic API.
"""
from dataclasses import dataclass
from typing import List, Union, cast
from enum import Enum
from anthropic import (
    Anthropic,
    AnthropicBedrock,
    AnthropicVertex,
    APIError,
    APIResponseValidationError,
    APIStatusError,
)
from anthropic.types.beta import (
    BetaCacheControlEphemeralParam,
    BetaContentBlockParam,
    BetaImageBlockParam,
    BetaMessage,
    BetaMessageParam,
    BetaTextBlock,
    BetaTextBlockParam,
    BetaToolResultBlockParam,
    BetaToolUseBlockParam,
)
from datetime import datetime

from .tools import ToolResult


COMPUTER_USE_BETA_FLAG_20250124 = "computer-use-2025-01-24"
COMPUTER_USE_BETA_FLAG_20251124 = "computer-use-2025-11-24"
COMPUTER_USE_BETA_FLAG = COMPUTER_USE_BETA_FLAG_20251124
PROMPT_CACHING_BETA_FLAG = "prompt-caching-2024-07-31"
COMPUTER_USE_TYPE_20250124 = "computer_20250124"
COMPUTER_USE_TYPE_20251124 = "computer_20251124"
COMPUTER_USE_TYPE = COMPUTER_USE_TYPE_20251124

# computer_20250124 computer-use-2025-01-24
# computer_20251124 computer-use-2025-11-24


@dataclass(frozen=True)
class ClaudeRuntimeProfile:
    """Version-specific Anthropic API behavior kept for reproducible runs."""

    label: str
    computer_use_beta_flag: str
    computer_use_type: str
    thinking_mode: str
    default_max_tokens: int
    use_claude_47_prompt: bool = False


CLAUDE_45_PROFILE = ClaudeRuntimeProfile(
    label="Claude 4.5 legacy implementation",
    computer_use_beta_flag=COMPUTER_USE_BETA_FLAG_20250124,
    computer_use_type=COMPUTER_USE_TYPE_20250124,
    thinking_mode="legacy",
    default_max_tokens=3000,
)
CLAUDE_46_PROFILE = ClaudeRuntimeProfile(
    label="Claude 4.6 legacy implementation",
    computer_use_beta_flag=COMPUTER_USE_BETA_FLAG_20251124,
    computer_use_type=COMPUTER_USE_TYPE_20251124,
    thinking_mode="legacy",
    default_max_tokens=3000,
)
CLAUDE_47_PROFILE = ClaudeRuntimeProfile(
    label="Claude 4.7 adaptive-effort implementation",
    computer_use_beta_flag=COMPUTER_USE_BETA_FLAG_20251124,
    computer_use_type=COMPUTER_USE_TYPE_20251124,
    thinking_mode="adaptive",
    default_max_tokens=16000,
    use_claude_47_prompt=True,
)
DEFAULT_CLAUDE_PROFILE = CLAUDE_46_PROFILE


def get_claude_runtime_profile(model_name: str) -> ClaudeRuntimeProfile:
    """Return the API/runtime profile for a Claude model name or provider model id."""
    normalized_model = (model_name or "").lower().replace("_", "-")
    if "4-7" in normalized_model or "4.7" in normalized_model:
        return CLAUDE_47_PROFILE
    if "4-6" in normalized_model or "4.6" in normalized_model:
        return CLAUDE_46_PROFILE
    if "4-5" in normalized_model or "4.5" in normalized_model:
        return CLAUDE_45_PROFILE
    return DEFAULT_CLAUDE_PROFILE


class APIProvider(Enum):
    ANTHROPIC = "anthropic"
    BEDROCK = "bedrock"
    OPENROUTER = "openrouter"
    VERTEX = "vertex"


PROVIDER_TO_DEFAULT_MODEL_NAME: dict[(APIProvider, str), str] = {
    (APIProvider.ANTHROPIC, "claude-3-5-sonnet-20241022"): "claude-3-5-sonnet-20241022",
    (APIProvider.BEDROCK, "claude-3-5-sonnet-20241022"): "us.anthropic.claude-3-5-sonnet-20241022-v2:0",
    (APIProvider.VERTEX, "claude-3-5-sonnet-20241022"): "claude-3-5-sonnet-v1@20241022",
    (APIProvider.ANTHROPIC, "claude-3-7-sonnet-20250219"): "claude-3-7-sonnet-20250219",
    (APIProvider.BEDROCK, "claude-3-7-sonnet-20250219"): "us.anthropic.claude-3-7-sonnet-20250219-v1:0",
    (APIProvider.VERTEX, "claude-3-7-sonnet-20250219"): "claude-3-7-sonnet-v1@20250219",
    (APIProvider.ANTHROPIC, "claude-4-opus-20250514"): "claude-4-opus-20250514",
    (APIProvider.BEDROCK, "claude-4-opus-20250514"): "us.anthropic.claude-opus-4-20250514-v1:0",
    (APIProvider.VERTEX, "claude-4-opus-20250514"): "claude-4-opus-v1@20250514",
    (APIProvider.ANTHROPIC, "claude-opus-4-20250514"): "claude-opus-4-20250514",
    (APIProvider.ANTHROPIC, "claude-opus-4-1-20250805"): "claude-opus-4-1-20250805",
    (APIProvider.ANTHROPIC, "claude-4-sonnet-20250514"): "claude-4-sonnet-20250514",
    (APIProvider.ANTHROPIC, "claude-sonnet-4-20250514"): "claude-sonnet-4-20250514",
    (APIProvider.BEDROCK, "claude-4-sonnet-20250514"): "us.anthropic.claude-sonnet-4-20250514-v1:0",
    (APIProvider.VERTEX, "claude-4-sonnet-20250514"): "claude-sonnet-4-v1@20250514",
    (APIProvider.ANTHROPIC, "claude-sonnet-4-5-20250929"): "claude-sonnet-4-5-20250929",
    (APIProvider.BEDROCK, "claude-sonnet-4-5-20250929"): "us.anthropic.claude-sonnet-4-5-20250929-v1:0",
    (APIProvider.VERTEX, "claude-sonnet-4-5-20250929"): "claude-sonnet-4-v1@20250929",
    (APIProvider.ANTHROPIC, "claude-opus-4-5"): "claude-opus-4-5-20251101",
    (APIProvider.BEDROCK, "claude-opus-4-5"): "us.anthropic.claude-opus-4-5-20251101-v1:0",
    (APIProvider.ANTHROPIC, "claude-opus-4-6"): "claude-opus-4-6",
    (APIProvider.BEDROCK, "claude-opus-4-6"): "global.anthropic.claude-opus-4-6-v1",
    (APIProvider.ANTHROPIC, "claude-opus-4-7"): "claude-opus-4-7",
    (APIProvider.BEDROCK, "claude-opus-4-7"): "global.anthropic.claude-opus-4-7-v1",
    (APIProvider.ANTHROPIC, "claude-sonnet-4-7"): "claude-sonnet-4-7",
    (APIProvider.BEDROCK, "claude-sonnet-4-7"): "global.anthropic.claude-sonnet-4-7-v1",
    (APIProvider.OPENROUTER, "claude-sonnet-4-20250514"): "anthropic/claude-sonnet-4",
    (APIProvider.OPENROUTER, "claude-sonnet-4-5-20250929"): "anthropic/claude-sonnet-4.5",
    (APIProvider.OPENROUTER, "claude-sonnet-4-6"): "anthropic/claude-sonnet-4.6",
    (APIProvider.OPENROUTER, "us.anthropic.claude-sonnet-4-6"): "anthropic/claude-sonnet-4.6",
    (APIProvider.OPENROUTER, "global.anthropic.claude-sonnet-4-6-v1"): "anthropic/claude-sonnet-4.6",
    (APIProvider.OPENROUTER, "claude-opus-4-20250514"): "anthropic/claude-opus-4",
    (APIProvider.OPENROUTER, "claude-opus-4-1-20250805"): "anthropic/claude-opus-4.1",
    (APIProvider.OPENROUTER, "claude-opus-4-5"): "anthropic/claude-opus-4.5",
    (APIProvider.OPENROUTER, "claude-opus-4-6"): "anthropic/claude-opus-4.6",
    (APIProvider.OPENROUTER, "global.anthropic.claude-opus-4-6-v1"): "anthropic/claude-opus-4.6",
    (APIProvider.OPENROUTER, "claude-opus-4-7"): "anthropic/claude-opus-4.7",
    (APIProvider.OPENROUTER, "global.anthropic.claude-opus-4-7-v1"): "anthropic/claude-opus-4.7",
}


def get_model_name(provider: APIProvider, model_name: str) -> str:
    """
    Get the actual model name to use for API calls.
    
    Simply returns the model name as-is for direct API usage.
    """
    # Look up in the mapping table
    key = (provider, model_name)
    if key in PROVIDER_TO_DEFAULT_MODEL_NAME:
        return PROVIDER_TO_DEFAULT_MODEL_NAME[key]

    if provider == APIProvider.BEDROCK and model_name.startswith(
        ("anthropic.", "us.anthropic.", "global.anthropic.")
    ):
        return model_name
    if provider == APIProvider.VERTEX and "@" in model_name:
        return model_name
    if provider == APIProvider.OPENROUTER:
        if model_name.startswith("anthropic/"):
            return model_name

        normalized_model = model_name
        for prefix in ("global.anthropic.", "us.anthropic.", "anthropic."):
            if normalized_model.startswith(prefix):
                normalized_model = normalized_model[len(prefix):]
                break
        for suffix in ("-v2:0", "-v1:0", "-v2", "-v1"):
            if normalized_model.endswith(suffix):
                normalized_model = normalized_model[: -len(suffix)]
                break
        normalized_model = normalized_model.replace("claude-sonnet-4-5", "claude-sonnet-4.5")
        normalized_model = normalized_model.replace("claude-sonnet-4-6", "claude-sonnet-4.6")
        normalized_model = normalized_model.replace("claude-sonnet-4-7", "claude-sonnet-4.7")
        normalized_model = normalized_model.replace("claude-opus-4-1", "claude-opus-4.1")
        normalized_model = normalized_model.replace("claude-opus-4-5", "claude-opus-4.5")
        normalized_model = normalized_model.replace("claude-opus-4-6", "claude-opus-4.6")
        normalized_model = normalized_model.replace("claude-opus-4-7", "claude-opus-4.7")
        if normalized_model.startswith("claude-"):
            return f"anthropic/{normalized_model}"
    
    # If not found and using Anthropic direct API, return as-is (may work for new models)
    if provider == APIProvider.ANTHROPIC:
        return model_name
    
    # For Bedrock/Vertex/OpenRouter, show available models if not found
    available_models = [k[1] for k in PROVIDER_TO_DEFAULT_MODEL_NAME.keys() if k[0] == provider]
    raise ValueError(
        f"❌ Model '{model_name}' is not configured for {provider.value}.\n"
        f"📋 Available models for {provider.value}: {available_models}"
    )


# This system prompt is optimized for the Docker environment in this repository and
# specific tool combinations enabled.
# We encourage modifying this system prompt to ensure the model has context for the
# environment it is running in, and to provide any additional information that may be
# helpful for the task at hand.
SYSTEM_PROMPT = f"""<SYSTEM_CAPABILITY>
* You are utilising an Ubuntu virtual machine using x86_64 architecture with internet access.
* You can feel free to install Ubuntu applications with your bash tool. Use curl instead of wget.
* To open browser, please just click on the Chrome icon.  Note, Chrome is what is installed on your system.
* Using bash tool you can start GUI applications, but you need to set export DISPLAY=:1 and use a subshell. For example "(DISPLAY=:1 xterm &)". GUI apps run with bash tool will appear within your desktop environment, but they may take some time to appear. Take a screenshot to confirm it did.
* When using your bash tool with commands that are expected to output very large quantities of text, redirect into a tmp file and use str_replace_editor or `grep -n -B <lines before> -A <lines after> <query> <filename>` to confirm output.
* When viewing a page it can be helpful to zoom out so that you can see everything on the page.  Either that, or make sure you scroll down to see everything before deciding something isn't available.
* DO NOT ask users for clarification during task execution. DO NOT stop to request more information from users. Always take action using available tools.
* When using your computer function calls, they take a while to run and send back to you.  Where possible/feasible, try to chain multiple of these calls all into one function calls request.
* TASK FEASIBILITY: You can declare a task infeasible at any point during execution - whether at the beginning after taking a screenshot, or later after attempting some actions and discovering barriers. Carefully evaluate whether the task is feasible given the current system state, available applications, and task requirements. If you determine that a task cannot be completed due to:
  - Missing required applications or dependencies that cannot be installed
  - Insufficient permissions or system limitations
  - Contradictory or impossible requirements
  - Any other fundamental barriers that make completion impossible
  Then you MUST output exactly "[INFEASIBLE]" (including the square brackets) anywhere in your response to trigger the fail action. The system will automatically detect this pattern and terminate the task appropriately.
* The current date is {datetime.today().strftime('%A, %B %d, %Y')}.
* Home directory of this Ubuntu system is '/home/user'.
* If you need a password for sudo, the password of the computer is 'osworld-public-evaluation'. 
</SYSTEM_CAPABILITY>

<IMPORTANT>
* If the item you are looking at is a pdf, if after taking a single screenshot of the pdf it seems that you want to read the entire document instead of trying to continue to read the pdf from your screenshots + navigation, determine the URL, use curl to download the pdf, install and use pdftotext to convert it to a text file, and then read that text file directly with your StrReplaceEditTool.
</IMPORTANT>"""

SYSTEM_PROMPT_WINDOWS = f"""<SYSTEM_CAPABILITY>
* You are utilising a Windows virtual machine using x86_64 architecture with internet access.
* To open browser, please just click on the Chrome icon.  Note, Chrome is what is installed on your system.
* When viewing a page it can be helpful to zoom out so that you can see everything on the page.  Either that, or make sure you scroll down to see everything before deciding something isn't available.
* The current date is {datetime.today().strftime('%A, %B %d, %Y')}.
* Home directory of this Windows system is 'C:\\Users\\user'.
* When you want to open some applications on Windows, please use Double Click on it instead of clicking once.
* If you need a password for sudo, The password of the computer is 'osworld-public-evaluation'.
</SYSTEM_CAPABILITY>"""


CLAUDE_47_PROMPT_ADDITIONS = """<CLAUDE_4_7_NOTES>
* DO NOT use LibreOffice macros or GIMP Script-Fu to complete tasks. Always use the GUI interface directly with mouse and keyboard actions. Macros and scripting cause reliability issues and task failures.
* For GIMP tasks, do NOT save or export files unless the instruction explicitly asks you to. Note that existing tasks that require file output will ask you to "export", not "save". Most GIMP tasks are evaluated automatically without requiring you to save.
* Before starting a task, consider whether it is achievable with the designated application's native GUI features. If the app fundamentally lacks the requested capability, declare it infeasible instead of using CLI tools, Python scripts, or other applications as workarounds.
* After completing a task, verify the visible or functional result. If your actions had no real effect, reconsider whether the task is feasible.
</CLAUDE_4_7_NOTES>"""


BATCHED_ACTION_PROMPT = """CRITICAL - ALWAYS BATCH ACTIONS: You MUST combine as many actions as possible into a SINGLE tool call.
Each tool call has significant latency overhead, so batching is essential for efficiency.
Do NOT make separate tool calls for actions you can predict the outcome of.

You do NOT need a screenshot after every action. Only take a screenshot when you genuinely
need to observe the result to decide what to do next. Many action sequences are predictable.

Example of GOOD batching (1 tool call instead of 6):
  actions=[
    {"action": "left_click", "coordinate": [300, 200]},
    {"action": "wait", "duration": 1},
    {"action": "triple_click", "coordinate": [400, 150]},
    {"action": "type", "text": "new value"},
    {"action": "key", "text": "Tab"},
    {"action": "type", "text": "another value"}
  ]

WRONG approach: making a separate tool call for each individual click, type, or keypress.
This wastes round trips and slows down task completion dramatically.

Think of each tool call as expensive. Ask yourself: "Can I predict what will happen?"
If yes, batch the action with subsequent actions. Only screenshot when the next step
depends on seeing the screen.

For form filling, menu navigation, and keyboard shortcuts, you almost never need a
screenshot between individual actions - batch them all together."""


BATCHED_TOOL_SCHEMA = {
    "name": "computer",
    "description": (
        "Execute one or more computer use actions in sequence. Batch multiple actions "
        "into a single call to minimize round trips. Actions use the same parameters as "
        "the Anthropic computer use tool: screenshot, left_click, right_click, middle_click, "
        "double_click, triple_click, type, key, scroll, mouse_move, left_click_drag, "
        "left_mouse_down, left_mouse_up, hold_key, and wait."
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "actions": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "action": {
                            "type": "string",
                            "enum": [
                                "key", "type", "mouse_move", "left_click",
                                "left_click_drag", "right_click", "middle_click",
                                "double_click", "triple_click", "scroll",
                                "hold_key", "screenshot", "cursor_position",
                                "left_mouse_down", "left_mouse_up", "wait",
                            ],
                            "description": "The action to perform. Same action types as the Anthropic computer use tool.",
                        },
                        "text": {
                            "type": "string",
                            "description": "For type: text to type. For key: key combo, e.g. 'ctrl+s'. For click/scroll: modifier key to hold during action.",
                        },
                        "coordinate": {
                            "type": "array",
                            "items": {"type": "integer"},
                            "minItems": 2,
                            "maxItems": 2,
                            "description": "[x, y] pixel coordinates on screen.",
                        },
                        "start_coordinate": {
                            "type": "array",
                            "items": {"type": "integer"},
                            "minItems": 2,
                            "maxItems": 2,
                            "description": "[x, y] start coordinates for left_click_drag.",
                        },
                        "scroll_direction": {
                            "type": "string",
                            "enum": ["up", "down", "left", "right"],
                            "description": "Direction to scroll. Required for action=scroll.",
                        },
                        "scroll_amount": {
                            "type": "integer",
                            "description": "Number of scroll clicks. Required for action=scroll.",
                        },
                        "duration": {
                            "type": "number",
                            "description": "Duration in seconds for hold_key or wait.",
                        },
                        "repeat": {
                            "type": "integer",
                            "minimum": 1,
                            "description": "Times to repeat key press. For action=key only.",
                        },
                    },
                    "required": ["action"],
                },
                "description": "List of actions to execute in sequence. Batch as many as possible.",
            }
        },
        "required": ["actions"],
    },
}



def _make_api_tool_result(
        result: ToolResult, tool_use_id: str
    ) -> BetaToolResultBlockParam:
        """Convert an agent ToolResult to an API ToolResultBlockParam."""
        tool_result_content: Union[List[Union[BetaTextBlockParam,
                                              BetaImageBlockParam]], str] = []
        is_error = False

        if not result or (result.get('error') is not None and result.get('error') != ""):
            is_error = True
            error_message = str(result.get('error', 'Unknown error occurred')) if result else 'No result received'
            tool_result_content = [{
                "type": "text",
                "text": _maybe_prepend_system_tool_result(result, error_message)
            }]
            
        else:
            if result.get('output'):
                tool_result_content.append({
                    "type": "text",
                    "text": _maybe_prepend_system_tool_result(
                        result,
                        str(result.get('output', '')
                            if result else '') 
                    ),
                })

            if result.get('base64_image'):
                tool_result_content.append({
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": "image/png",
                        "data": result.get('base64_image', ''),
                    },
                })
            
            if not tool_result_content:
                tool_result_content.append({
                    "type": "text",
                    "text": "Action completed successfully"
                })

        return {
            "type": "tool_result",
            "content": tool_result_content,
            "tool_use_id": tool_use_id,
            "is_error": is_error,
        }

def _maybe_prepend_system_tool_result(result: ToolResult, result_text: str):
    if not result:
        return result_text

    if result.get('system', False):
        result_text = f"<system>{result.get('system','')}</system>\n{result_text}"
    return result_text



def _inject_prompt_caching(
    messages: list[BetaMessageParam],
):
    """
    Set cache breakpoints for the 3 most recent turns
    one cache breakpoint is left for tools/system prompt, to be shared across sessions
    """

    breakpoints_remaining = 2  # Use full budget for recent messages
    messages_processed = 0
    
    for message in reversed(messages):
        if message["role"] == "user" and isinstance(
            content := message["content"], list
        ):
            messages_processed += 1
            # Check if this message would fit within the remaining budget
            if breakpoints_remaining >= len(content):
                # We have enough budget, spend it and add cache_control
                breakpoints_remaining -= len(content)
                # Use type ignore to bypass TypedDict check until SDK types are updated
                content[-1]["cache_control"] = BetaCacheControlEphemeralParam(  # type: ignore
                    {"type": "ephemeral"}
                )
            else:
                # Check if this is the first message (contains image + text with task description)
                is_first_message = messages_processed == len([msg for msg in messages if msg["role"] == "user"])
                
                if not is_first_message:
                    # Not enough budget, remove any existing cache_control from this message
                    content[-1].pop("cache_control", None)
                # Continue to clean up older messages that might have cache_control from previous turns


def _maybe_filter_to_n_most_recent_images(
    messages: list[BetaMessageParam],
    images_to_keep: int,
    min_removal_threshold: int,
):
    """
    With the assumption that images are screenshots that are of diminishing value as
    the conversation progresses, remove all but the final `images_to_keep` tool_result
    images in place, with a chunk of min_removal_threshold to reduce the amount we
    break the implicit prompt cache.
    """
    if images_to_keep is None:
        return messages

    tool_result_blocks = cast(
        list[BetaToolResultBlockParam],
        [
            item
            for message in messages
            for item in (
                message["content"] if isinstance(message["content"], list) else []
            )
            if isinstance(item, dict) and item.get("type") == "tool_result"
        ],
    )

    total_images = sum(
        1
        for tool_result in tool_result_blocks
        for content in tool_result.get("content", [])
        if isinstance(content, dict) and content.get("type") == "image"
    )

    images_to_remove = total_images - images_to_keep
    # for better cache behavior, we want to remove in chunks
    images_to_remove -= images_to_remove % min_removal_threshold

    for tool_result in tool_result_blocks:
        if isinstance(tool_result.get("content"), list):
            new_content = []
            for content in tool_result.get("content", []):
                if isinstance(content, dict) and content.get("type") == "image":
                    if images_to_remove > 0:
                        images_to_remove -= 1
                        continue
                new_content.append(content)
            tool_result["content"] = new_content


def validate_model_support(
    model_name: str,
    provider: APIProvider = APIProvider.BEDROCK,
    api_key: str = None,
    temperature: float = None,
    top_p: float = None,
    no_thinking: bool = False,
    use_isp: bool = False,
    effort: str = "max",
) -> bool:
    """
    Validate model support with the same API call pattern as the main agent.
    
    Args:
        model_name: The model name to validate
        provider: API provider (ANTHROPIC, BEDROCK, or VERTEX)
        api_key: Optional API key (only for ANTHROPIC provider), defaults to ANTHROPIC_API_KEY env var
        temperature: Optional temperature parameter for testing
        top_p: Optional top_p parameter for testing
        no_thinking: Disable thinking mode (matches AnthropicAgent)
        use_isp: Use interleaved scratchpad mode for Claude 4.5/4.6 legacy runs
        effort: Effort level for Claude 4.7 adaptive thinking
        
    Returns:
        True if model is supported and API call succeeds, False otherwise
    """
    profile = get_claude_runtime_profile(model_name)
    print(f"🔍 Validating model support: {model_name} (provider: {provider.value}, profile: {profile.label})")
    
    try:
        from anthropic import Anthropic, AnthropicBedrock, AnthropicVertex
        import os
        import time
        
        # Same client setup as main agent - choose based on provider
        if provider == APIProvider.ANTHROPIC:
            client = Anthropic(
                api_key=api_key or os.environ.get("ANTHROPIC_API_KEY"), 
                max_retries=4
            ).with_options(default_headers={"anthropic-beta": profile.computer_use_beta_flag})
        elif provider == APIProvider.OPENROUTER:
            openrouter_key = (
                api_key
                or os.environ.get("OPENROUTER_API_KEY")
                or os.environ.get("OPENROUTER_KEY")
            )
            if not openrouter_key:
                print("❌ OPENROUTER_API_KEY or OPENROUTER_KEY must be set for OpenRouter")
                return False
            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",
                },
            )
        elif provider == APIProvider.BEDROCK:
            client = AnthropicBedrock(
                aws_access_key=os.getenv('AWS_ACCESS_KEY_ID'),
                aws_secret_key=os.getenv('AWS_SECRET_ACCESS_KEY'),
                aws_region=os.getenv('AWS_DEFAULT_REGION'),
            )
        elif provider == APIProvider.VERTEX:
            client = AnthropicVertex()
        
        # Same message format as main agent - always use structured format with cache_control
        messages = [{"role": "user", "content": [{"type": "text", "text": "Respond with 'OK'", "cache_control": {"type": "ephemeral"}}]}]
        
        # Same betas configuration as main agent
        betas = [profile.computer_use_beta_flag]
        if use_isp and profile.thinking_mode == "legacy":
            betas.append("interleaved-thinking-2025-05-14")
            
        system = [{"type": "text", "text": "You are Claude. Respond with 'OK'."}]
        
        # Same tools configuration as main agent, selected by provider/version.
        if provider == APIProvider.OPENROUTER:
            tools = [BATCHED_TOOL_SCHEMA]
        else:
            tools = [{
                "name": "computer",
                "type": profile.computer_use_type,
                "display_width_px": 1280,
                "display_height_px": 720,
                "display_number": 1,
            }]
        
        # Same thinking configuration as main agent
        max_tokens = 50  # Base validation max_tokens
        if profile.thinking_mode == "adaptive":
            extra_body = {
                "thinking": {"type": "adaptive"},
                "output_config": {"effort": effort},
            }
            actual_max_tokens = profile.default_max_tokens
        else:
            if no_thinking:
                extra_body = {}
                actual_max_tokens = max_tokens
            else:
                budget_tokens = 2048
                # Same logic as main agent: if max_tokens <= budget_tokens, increase it
                if max_tokens <= budget_tokens:
                    actual_max_tokens = budget_tokens + 500
                else:
                    actual_max_tokens = max_tokens
                extra_body = {
                    "thinking": {"type": "enabled", "budget_tokens": budget_tokens}
                }
        
        # Sampling parameters (same logic as main agent)
        sampling_params = {}
        if temperature is not None:
            sampling_params['temperature'] = temperature
        if top_p is not None:
            sampling_params['top_p'] = top_p
        
        # Retry logic with 5 attempts, 5 second delays
        for attempt in range(5):
            try:
                # Same API call pattern as main agent
                create_kwargs = {
                    "max_tokens": actual_max_tokens,
                    "messages": messages,
                    "model": get_model_name(provider, model_name),
                    "system": system,
                    "tools": tools,
                    "betas": betas,
                    "extra_body": extra_body,
                    **sampling_params,
                }
                if provider == APIProvider.OPENROUTER:
                    create_kwargs["cache_control"] = {"type": "ephemeral"}
                    create_kwargs["extra_body"] = {
                        **extra_body,
                        "session_id": "osworld-openrouter-validation",
                    }

                client.beta.messages.create(
                    **create_kwargs
                )
                
                print(f"✅ Model {model_name} validated successfully with {provider.value}")
                return True
            except Exception as e:
                error_msg = str(e)
                if attempt < 4:  # Don't print error on final attempt
                    print(f"🔄 Validation attempt {attempt + 1}/5 failed: \"{error_msg}\"")
                    print(f"⏳ Retrying in 5 seconds...")
                    time.sleep(5)
                else:
                    print(f"❌ All validation attempts failed. Final error: \"{error_msg}\"")
        
        return False
        
    except ValueError:
        return False
    except Exception as e:
        print(f"❌ API validation setup failed: {e}")
        return False


def _response_to_params(
    response: BetaMessage,
) -> list[BetaContentBlockParam]:
    res: list[BetaContentBlockParam] = []
    if response.content:
        for block in response.content:
            if isinstance(block, BetaTextBlock):
                if block.text:
                    res.append(BetaTextBlockParam(type="text", text=block.text))
                elif getattr(block, "type", None) == "thinking":
                    # Handle thinking blocks - include signature field
                    thinking_block = {
                        "type": "thinking",
                        "thinking": getattr(block, "thinking", None),
                    }
                    if hasattr(block, "signature"):
                        thinking_block["signature"] = getattr(block, "signature", None)
                    res.append(cast(BetaContentBlockParam, thinking_block))
            else:
                # Handle tool use blocks normally
                res.append(cast(BetaToolUseBlockParam, block.model_dump()))
        return res
    else:
        return []
