import re
from typing import Any

CODE_EXECUTOR_TOOL_DESCRIPTION = """Executes Python code in an IPython notebook cell:
* The output generated by the notebook cell is returned by this tool
* State is persistent across executions and discussions with the user
* The input code may reference variables created in previous executions
"""


CODE_EXECUTOR_TOOL = {
    "type": "function",
    "function": {
        "name": "execute_ipython_cell",
        "description": CODE_EXECUTOR_TOOL_DESCRIPTION,
        "parameters": {
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": "The Python code to execute. Supports magic commands like !pip e.g. to install missing Python libraries via `!pip install <library>`.",
                }
            },
            "required": ["code"],
        },
    },
}


CODE_EDITOR_DESCRIPTION = """Custom editing tool for viewing, creating and editing files
* State is persistent across command calls and discussions with the user
* Interpret `path` as a relative path to the current working directory
* If `path` is a file, `view` displays the result of applying `cat -n`. If `path` is a directory, `view` lists non-hidden files and directories up to 2 levels deep
* The `create` command cannot be used if the specified `path` already exists as a file
* If a `command` generates a long output, it will be truncated and marked with `<response clipped>`
* The `undo_edit` command will revert the last edit made to the file at `path`

Notes for using the `str_replace` command:
* The `old_str` parameter should match EXACTLY one or more consecutive lines from the original file. Be mindful of whitespaces!
* If the `old_str` parameter is not unique in the file, the replacement will not be performed. Make sure to include enough context in `old_str` to make it unique
* The `new_str` parameter should contain the edited lines that should replace the `old_str`
"""


CODE_EDITOR_TOOL = {
    "type": "function",
    "function": {
        "name": "str_replace_editor",
        "description": CODE_EDITOR_DESCRIPTION,
        "parameters": {
            "type": "object",
            "properties": {
                "command": {
                    "description": "The commands to run. Allowed options are: `view`, `create`, `str_replace`, `insert`, `undo_edit`.",
                    "enum": ["view", "create", "str_replace", "insert", "undo_edit"],
                    "type": "string",
                },
                "path": {
                    "description": "Relative path to file or directory, e.g. `my_dir/my_file.py` or `my_dir`.",
                    "type": "string",
                },
                "file_text": {
                    "description": "Required parameter of `create` command, with the content of the file to be created.",
                    "type": "string",
                },
                "old_str": {
                    "description": "Required parameter of `str_replace` command containing the string in `path` to replace.",
                    "type": "string",
                },
                "new_str": {
                    "description": "Optional parameter of `str_replace` command containing the new string (if not given, no string will be added). Required parameter of `insert` command containing the string to insert.",
                    "type": "string",
                },
                "insert_line": {
                    "description": "Required parameter of `insert` command. The `new_str` will be inserted AFTER the line `insert_line` of `path`.",
                    "type": "integer",
                },
                "view_range": {
                    "description": "Optional parameter of `view` command when `path` points to a file. If none is given, the full file is shown. If provided, the file will be shown in the indicated line number range, e.g. [11, 12] will show lines 11 and 12. Indexing at 1 to start. Setting `[start_line, -1]` shows all lines from `start_line` to the end of the file.",
                    "items": {"type": "integer"},
                    "type": "array",
                },
            },
            "required": ["command", "path"],
        },
    },
}


CLAUDE_3_7_SONNET_MODELS = [
    "anthropic/claude-3-7-sonnet-latest",
    "anthropic/claude-3-7-sonnet-20250219",
    "bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0",
    "vertex_ai/claude-3-7-sonnet@20250219",
]

CLAUDE_3_5_SONNET_MODELS = [
    "anthropic/claude-3-5-sonnet-latest",
    "anthropic/claude-3-5-sonnet-20241022",
    "bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0",
    "vertex_ai/claude-3-5-sonnet-v2@20241022",
]

CLAUDE_3_5_HAIKU_MODELS = [
    "anthropic/claude-3-5-haiku-latest",
    "anthropic/claude-3-5-haiku-20241022",
    "bedrock/anthropic.claude-3-5-haiku-20241022-v1:0",
    "vertex_ai/claude-3-5-haiku@20241022",
]

CLAUDE_MODELS = CLAUDE_3_7_SONNET_MODELS + CLAUDE_3_5_SONNET_MODELS + CLAUDE_3_5_HAIKU_MODELS


def beta_flag(model_name: str):
    if model_name in CLAUDE_3_7_SONNET_MODELS:
        return {"anthropic-beta": "computer-use-2025-01-24"}
    elif model_name in CLAUDE_3_5_SONNET_MODELS:
        return {"anthropic-beta": "computer-use-2024-10-22"}
    else:
        return {}


def code_editor_tool(model_name: str):
    if model_name in CLAUDE_3_7_SONNET_MODELS:
        return code_editor_tool_ref("text_editor_20250124")
    elif model_name in CLAUDE_3_5_SONNET_MODELS:
        return code_editor_tool_ref("text_editor_20241022")
    elif "gemini" in model_name:
        return CODE_EDITOR_TOOL["function"]
    else:
        return CODE_EDITOR_TOOL


def code_editor_tool_ref(editor_type: str):
    return {
        "type": f"{editor_type}",
        "function": {
            "name": "str_replace_editor",
        },
    }


def code_executor_tool(model_name: str):
    if "gemini" in model_name:
        return CODE_EXECUTOR_TOOL["function"]
    else:
        return CODE_EXECUTOR_TOOL


def tool_name(tool: dict[str, Any]) -> str:
    if "function" in tool:
        return tool["function"]["name"]
    else:
        return tool["name"]  # Gemini models


def sanitize_tool_name(tool_name: str) -> str:
    return re.sub(r"[^a-zA-Z0-9_-]", "_", tool_name)


def code_editor_tool_use_default(model_name: str, provider_name: str) -> bool:
    return model_name in CLAUDE_MODELS or provider_name == "openai"


def code_executor_tool_use_default(model_name: str, provider_name: str) -> bool:
    return model_name in CLAUDE_MODELS or provider_name == "openai"
