import json
from dataclasses import dataclass
from os import getenv
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple, Type, Union

from pydantic import BaseModel

from agno.exceptions import ModelProviderError
from agno.models.base import Model
from agno.models.message import Message
from agno.models.metrics import MessageMetrics
from agno.models.response import ModelResponse
from agno.run.agent import RunOutput
from agno.utils.log import log_debug, log_error, log_warning
from agno.utils.tokens import count_schema_tokens

try:
    from boto3 import client as AwsClient
    from boto3.session import Session
    from botocore.exceptions import ClientError
except ImportError:
    raise ImportError("`boto3` not installed. Please install using `pip install boto3`")

try:
    import aioboto3

    AIOBOTO3_AVAILABLE = True
except ImportError:
    aioboto3 = None
    AIOBOTO3_AVAILABLE = False


BEDROCK_SUPPORTED_IMAGE_FORMATS = ["png", "jpeg", "webp", "gif"]
BEDROCK_SUPPORTED_VIDEO_FORMATS = ["mp4", "mov", "mkv", "webm", "flv", "mpeg", "mpg", "wmv", "three_gp"]
BEDROCK_SUPPORTED_FILE_FORMATS = ["pdf", "csv", "doc", "docx", "xls", "xlsx", "html", "txt", "md"]


@dataclass
class AwsBedrock(Model):
    """
    AWS Bedrock model.

    To use this model, you need to either:
    1. Set the following environment variables:
       - AWS_ACCESS_KEY_ID
       - AWS_SECRET_ACCESS_KEY
       - AWS_REGION
    2. Or provide a boto3 Session object

    For async support, you also need aioboto3 installed:
       pip install aioboto3

    Not all Bedrock models support all features. See this documentation for more information: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-supported-models-features.html

    Args:
        aws_region (Optional[str]): The AWS region to use.
        aws_access_key_id (Optional[str]): The AWS access key ID to use.
        aws_secret_access_key (Optional[str]): The AWS secret access key to use.
        aws_sso_auth (Optional[str]): Removes the need for an access and secret access key by leveraging the current profile's authentication
        session (Optional[Session]): A boto3 Session object to use for authentication.
    """

    id: str = "mistral.mistral-small-2402-v1:0"
    name: str = "AwsBedrock"
    provider: str = "AwsBedrock"

    aws_sso_auth: Optional[bool] = False
    aws_region: Optional[str] = None
    aws_access_key_id: Optional[str] = None
    aws_secret_access_key: Optional[str] = None
    aws_session_token: Optional[str] = None
    session: Optional[Session] = None

    # Request parameters
    max_tokens: Optional[int] = None
    temperature: Optional[float] = None
    top_p: Optional[float] = None
    stop_sequences: Optional[List[str]] = None
    request_params: Optional[Dict[str, Any]] = None

    client: Optional[AwsClient] = None
    async_client: Optional[Any] = None
    async_session: Optional[Any] = None

    def get_client(self) -> AwsClient:
        """
        Get the Bedrock client.

        Returns:
            AwsClient: The Bedrock client.
        """
        # When using a boto3 session, always recreate the client so
        # session credentials can be refreshed (IAM roles, EKS, STS).
        if not self.session and self.client is not None:
            return self.client

        # Return directly (not via self.client) so concurrent callers
        # on the same model instance each get their own client.
        if self.session:
            return self.session.client(
                "bedrock-runtime",
                region_name=self.aws_region or self.session.region_name,
            )

        self.aws_access_key_id = self.aws_access_key_id or getenv("AWS_ACCESS_KEY_ID")
        self.aws_secret_access_key = self.aws_secret_access_key or getenv("AWS_SECRET_ACCESS_KEY")
        self.aws_session_token = self.aws_session_token or getenv("AWS_SESSION_TOKEN")
        self.aws_region = self.aws_region or getenv("AWS_REGION")

        if self.aws_sso_auth:
            self.client = AwsClient(service_name="bedrock-runtime", region_name=self.aws_region)
        else:
            if not self.aws_access_key_id or not self.aws_secret_access_key:
                log_error(
                    "AWS credentials not found. Please set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables or provide a boto3 session."
                )

            self.client = AwsClient(
                service_name="bedrock-runtime",
                region_name=self.aws_region,
                aws_access_key_id=self.aws_access_key_id,
                aws_secret_access_key=self.aws_secret_access_key,
                aws_session_token=self.aws_session_token,
            )
        return self.client

    def get_async_client(self):
        """
        Get the async Bedrock client context manager.

        Returns:
            The async Bedrock client context manager.
        """
        if not AIOBOTO3_AVAILABLE:
            raise ImportError(
                "`aioboto3` not installed. Please install using `pip install aioboto3` for async support."
            )

        # When using a boto3 session, create the aioboto3 session from it
        # so that session credentials (IAM roles, EKS, STS) are respected.
        if self.session:
            credentials = self.session.get_credentials()
            if credentials is None:
                raise ValueError(
                    "boto3 session has no credentials. Check your AWS configuration "
                    "(environment variables, config files, IAM role, etc.)."
                )
            # Use local variables (not self.async_session) so concurrent
            # callers each get their own session and client.
            frozen = credentials.get_frozen_credentials()
            async_session = aioboto3.Session(
                aws_access_key_id=frozen.access_key,
                aws_secret_access_key=frozen.secret_key,
                aws_session_token=frozen.token,
                region_name=self.aws_region or self.session.region_name,
            )
            return async_session.client("bedrock-runtime")

        if self.async_session is None:
            self.aws_access_key_id = self.aws_access_key_id or getenv("AWS_ACCESS_KEY_ID")
            self.aws_secret_access_key = self.aws_secret_access_key or getenv("AWS_SECRET_ACCESS_KEY")
            self.aws_session_token = self.aws_session_token or getenv("AWS_SESSION_TOKEN")
            self.aws_region = self.aws_region or getenv("AWS_REGION")

            self.async_session = aioboto3.Session()

        client_kwargs: Dict[str, Any] = {
            "service_name": "bedrock-runtime",
            "region_name": self.aws_region,
        }

        if self.aws_sso_auth:
            pass
        else:
            if not self.aws_access_key_id or not self.aws_secret_access_key:
                import os

                env_access_key = os.environ.get("AWS_ACCESS_KEY_ID")
                env_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
                env_session_token = os.environ.get("AWS_SESSION_TOKEN")
                env_region = os.environ.get("AWS_REGION")

                if env_access_key and env_secret_key:
                    self.aws_access_key_id = env_access_key
                    self.aws_secret_access_key = env_secret_key
                    self.aws_session_token = env_session_token
                    if env_region:
                        self.aws_region = env_region
                        client_kwargs["region_name"] = self.aws_region

            if self.aws_access_key_id and self.aws_secret_access_key:
                client_kwargs["aws_access_key_id"] = self.aws_access_key_id
                client_kwargs["aws_secret_access_key"] = self.aws_secret_access_key
                if self.aws_session_token:
                    client_kwargs["aws_session_token"] = self.aws_session_token

        return self.async_session.client(**client_kwargs)

    def to_dict(self) -> Dict[str, Any]:
        """
        Convert the model to a dictionary.

        Returns:
            Dict[str, Any]: The dictionary representation of the model.
        """
        model_dict = super().to_dict()
        model_dict.update(
            {
                "aws_region": self.aws_region,
                "max_tokens": self.max_tokens,
                "temperature": self.temperature,
                "top_p": self.top_p,
                "stop_sequences": self.stop_sequences,
            }
        )
        cleaned_dict = {k: v for k, v in model_dict.items() if v is not None}
        return cleaned_dict

    def _format_tools_for_request(self, tools: Optional[List[Dict[str, Any]]]) -> List[Dict[str, Any]]:
        """
        Format the tools for the request.

        Returns:
            List[Dict[str, Any]]: The formatted tools.
        """
        parsed_tools = []
        if tools is not None:
            for tool_def in tools:
                func_def = tool_def.get("function", {})
                properties = {}
                required = []

                for param_name, param_info in func_def.get("parameters", {}).get("properties", {}).items():
                    properties[param_name] = param_info.copy()

                    if "description" not in properties[param_name]:
                        properties[param_name]["description"] = ""

                    if "null" not in (
                        param_info.get("type") if isinstance(param_info.get("type"), list) else [param_info.get("type")]
                    ):
                        required.append(param_name)

                parsed_tools.append(
                    {
                        "toolSpec": {
                            "name": func_def.get("name") or "",
                            "description": func_def.get("description") or "",
                            "inputSchema": {"json": {"type": "object", "properties": properties, "required": required}},
                        }
                    }
                )

        return parsed_tools

    def _get_inference_config(self) -> Dict[str, Any]:
        """
        Get the inference config.

        Returns:
            Dict[str, Any]: The inference config.
        """
        request_kwargs = {
            "maxTokens": self.max_tokens,
            "temperature": self.temperature,
            "topP": self.top_p,
            "stopSequences": self.stop_sequences,
        }

        return {k: v for k, v in request_kwargs.items() if v is not None}

    def _format_messages(
        self, messages: List[Message], compress_tool_results: bool = False
    ) -> Tuple[List[Dict[str, Any]], Optional[List[Dict[str, Any]]]]:
        """
        Format the messages for the request.

        Args:
            messages: List of messages to format
            compress_tool_results: Whether to compress tool results

        Returns:
            Tuple[List[Dict[str, Any]], Optional[List[Dict[str, Any]]]]: The formatted messages.
        """
        from agno.utils.message import normalize_tool_messages

        # Backwards compat: expand old Gemini combined tool messages into individual canonical messages
        messages = normalize_tool_messages(messages)

        formatted_messages: List[Dict[str, Any]] = []
        system_message = None
        for message in messages:
            if message.role == "system":
                system_message = [{"text": message.content}]
            elif message.role == "tool":
                content = message.get_content(use_compressed_content=compress_tool_results)
                tool_result = {
                    "toolUseId": message.tool_call_id,
                    "content": [{"json": {"result": content}}],
                }
                if (
                    formatted_messages
                    and formatted_messages[-1]["role"] == "user"
                    and formatted_messages[-1]["content"]
                    and "toolResult" in formatted_messages[-1]["content"][0]
                ):
                    formatted_messages[-1]["content"].append({"toolResult": tool_result})
                else:
                    formatted_message: Dict[str, Any] = {"role": "user", "content": [{"toolResult": tool_result}]}
                    formatted_messages.append(formatted_message)
            else:
                formatted_message = {"role": message.role, "content": []}
                if isinstance(message.content, list):
                    formatted_message["content"].extend(message.content)
                elif message.tool_calls:
                    tool_use_content = []
                    for tool_call in message.tool_calls:
                        try:
                            # Parse arguments with error handling for empty or invalid JSON
                            arguments = tool_call["function"]["arguments"]
                            if not arguments or arguments.strip() == "":
                                tool_input = {}
                            else:
                                tool_input = json.loads(arguments)
                        except (json.JSONDecodeError, KeyError) as e:
                            log_warning(f"Failed to parse tool call arguments: {e}")
                            tool_input = {}

                        tool_use_content.append(
                            {
                                "toolUse": {
                                    "toolUseId": tool_call["id"],
                                    "name": tool_call["function"]["name"],
                                    "input": tool_input,
                                }
                            }
                        )
                    formatted_message["content"].extend(tool_use_content)
                else:
                    formatted_message["content"].append({"text": message.content})

                if message.images:
                    for image in message.images:
                        if not image.content:
                            raise ValueError("Image content is required for AWS Bedrock.")
                        if not image.format:
                            raise ValueError("Image format is required for AWS Bedrock.")

                        if image.format not in BEDROCK_SUPPORTED_IMAGE_FORMATS:
                            raise ValueError(
                                f"Unsupported image format: {image.format}. "
                                f"Supported formats: {BEDROCK_SUPPORTED_IMAGE_FORMATS}"
                            )

                        formatted_message["content"].append(
                            {
                                "image": {
                                    "format": image.format,
                                    "source": {
                                        "bytes": image.content,
                                    },
                                }
                            }
                        )
                if message.audio:
                    log_warning("Audio input is currently unsupported.")

                if message.videos:
                    for video in message.videos:
                        if not video.content:
                            raise ValueError("Video content is required for AWS Bedrock.")
                        if not video.format:
                            raise ValueError("Video format is required for AWS Bedrock.")

                        if video.format not in BEDROCK_SUPPORTED_VIDEO_FORMATS:
                            raise ValueError(
                                f"Unsupported video format: {video.format}. "
                                f"Supported formats: {BEDROCK_SUPPORTED_VIDEO_FORMATS}"
                            )

                        formatted_message["content"].append(
                            {
                                "video": {
                                    "format": video.format,
                                    "source": {
                                        "bytes": video.content,
                                    },
                                }
                            }
                        )

                if message.files:
                    for file in message.files:
                        if not file.content:
                            raise ValueError("File content is required for AWS Bedrock document input.")
                        if not file.format:
                            raise ValueError("File format is required for AWS Bedrock document input.")
                        if not file.name:
                            raise ValueError("File name is required for AWS Bedrock document input.")

                        if file.format not in BEDROCK_SUPPORTED_FILE_FORMATS:
                            raise ValueError(
                                f"Unsupported file format: {file.format}. "
                                f"Supported formats: {BEDROCK_SUPPORTED_FILE_FORMATS}"
                            )

                        formatted_message["content"].append(
                            {
                                "document": {
                                    "format": file.format,
                                    "name": file.name,
                                    "source": {
                                        "bytes": file.content,
                                    },
                                }
                            }
                        )

                formatted_messages.append(formatted_message)
        # TODO: Add caching: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-call.html
        return formatted_messages, system_message

    def count_tokens(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_schema: Optional[Union[Dict, Type[BaseModel]]] = None,
    ) -> int:
        try:
            formatted_messages, system_message = self._format_messages(messages, compress_tool_results=True)
            converse_input: Dict[str, Any] = {"messages": formatted_messages}
            if system_message:
                converse_input["system"] = system_message

            response = self.get_client().count_tokens(modelId=self.id, input={"converse": converse_input})
            tokens = response.get("inputTokens", 0)

            # Count tool tokens
            if tools:
                from agno.utils.tokens import count_tool_tokens

                tokens += count_tool_tokens(tools, self.id)

            # Count schema tokens
            tokens += count_schema_tokens(output_schema, self.id)

            return tokens
        except Exception as e:
            log_warning(f"Failed to count tokens via Bedrock API: {e}")
            return super().count_tokens(messages, tools, output_schema)

    async def acount_tokens(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_schema: Optional[Union[Dict, Type[BaseModel]]] = None,
    ) -> int:
        try:
            formatted_messages, system_message = self._format_messages(messages, compress_tool_results=True)
            converse_input: Dict[str, Any] = {"messages": formatted_messages}
            if system_message:
                converse_input["system"] = system_message

            async with self.get_async_client() as client:
                response = await client.count_tokens(modelId=self.id, input={"converse": converse_input})
            tokens = response.get("inputTokens", 0)

            # Count tool tokens
            if tools:
                from agno.utils.tokens import count_tool_tokens

                tokens += count_tool_tokens(tools, self.id)

            # Count schema tokens
            tokens += count_schema_tokens(output_schema, self.id)

            return tokens
        except Exception as e:
            log_warning(f"Failed to count tokens via Bedrock API: {e}")
            return await super().acount_tokens(messages, tools, output_schema)

    def invoke(
        self,
        messages: List[Message],
        assistant_message: Message,
        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
        tools: Optional[List[Dict[str, Any]]] = None,
        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
        run_response: Optional[RunOutput] = None,
        compress_tool_results: bool = False,
    ) -> ModelResponse:
        """
        Invoke the Bedrock API.
        """
        try:
            formatted_messages, system_message = self._format_messages(messages, compress_tool_results)

            tool_config = None
            if tools:
                tool_config = {"tools": self._format_tools_for_request(tools)}

            body = {
                "system": system_message,
                "toolConfig": tool_config,
                "inferenceConfig": self._get_inference_config(),
            }
            body = {k: v for k, v in body.items() if v is not None}

            if self.request_params:
                log_debug(f"Calling {self.provider} with request parameters: {self.request_params}", log_level=2)
                body.update(**self.request_params)

            assistant_message.metrics.start_timer()
            response = self.get_client().converse(modelId=self.id, messages=formatted_messages, **body)
            assistant_message.metrics.stop_timer()

            model_response = self._parse_provider_response(response, response_format=response_format)

            return model_response

        except ClientError as e:
            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
            raise ModelProviderError(message=str(e.response), model_name=self.name, model_id=self.id) from e
        except Exception as e:
            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e

    def invoke_stream(
        self,
        messages: List[Message],
        assistant_message: Message,
        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
        tools: Optional[List[Dict[str, Any]]] = None,
        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
        run_response: Optional[RunOutput] = None,
        compress_tool_results: bool = False,
    ) -> Iterator[ModelResponse]:
        """
        Invoke the Bedrock API with streaming.
        """
        try:
            formatted_messages, system_message = self._format_messages(messages, compress_tool_results)

            tool_config = None
            if tools:
                tool_config = {"tools": self._format_tools_for_request(tools)}

            body = {
                "system": system_message,
                "toolConfig": tool_config,
                "inferenceConfig": self._get_inference_config(),
            }
            body = {k: v for k, v in body.items() if v is not None}

            if self.request_params:
                body.update(**self.request_params)

            assistant_message.metrics.start_timer()

            # Track current tool being built across chunks
            current_tool: Dict[str, Any] = {}

            for chunk in self.get_client().converse_stream(modelId=self.id, messages=formatted_messages, **body)[
                "stream"
            ]:
                model_response, current_tool = self._parse_provider_response_delta(chunk, current_tool)
                yield model_response

            assistant_message.metrics.stop_timer()

        except ClientError as e:
            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
            raise ModelProviderError(message=str(e.response), model_name=self.name, model_id=self.id) from e
        except Exception as e:
            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e

    async def ainvoke(
        self,
        messages: List[Message],
        assistant_message: Message,
        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
        tools: Optional[List[Dict[str, Any]]] = None,
        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
        run_response: Optional[RunOutput] = None,
        compress_tool_results: bool = False,
    ) -> ModelResponse:
        """
        Async invoke the Bedrock API.
        """
        try:
            formatted_messages, system_message = self._format_messages(messages, compress_tool_results)

            tool_config = None
            if tools:
                tool_config = {"tools": self._format_tools_for_request(tools)}

            body = {
                "system": system_message,
                "toolConfig": tool_config,
                "inferenceConfig": self._get_inference_config(),
            }
            body = {k: v for k, v in body.items() if v is not None}

            if self.request_params:
                log_debug(f"Calling {self.provider} with request parameters: {self.request_params}", log_level=2)
                body.update(**self.request_params)

            assistant_message.metrics.start_timer()

            async with self.get_async_client() as client:
                response = await client.converse(modelId=self.id, messages=formatted_messages, **body)

            assistant_message.metrics.stop_timer()

            model_response = self._parse_provider_response(response, response_format=response_format)

            return model_response

        except ClientError as e:
            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
            raise ModelProviderError(message=str(e.response), model_name=self.name, model_id=self.id) from e
        except Exception as e:
            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e

    async def ainvoke_stream(
        self,
        messages: List[Message],
        assistant_message: Message,
        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
        tools: Optional[List[Dict[str, Any]]] = None,
        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
        run_response: Optional[RunOutput] = None,
        compress_tool_results: bool = False,
    ) -> AsyncIterator[ModelResponse]:
        """
        Async invoke the Bedrock API with streaming.
        """
        try:
            formatted_messages, system_message = self._format_messages(messages, compress_tool_results)

            tool_config = None
            if tools:
                tool_config = {"tools": self._format_tools_for_request(tools)}

            body = {
                "system": system_message,
                "toolConfig": tool_config,
                "inferenceConfig": self._get_inference_config(),
            }
            body = {k: v for k, v in body.items() if v is not None}

            if self.request_params:
                body.update(**self.request_params)

            assistant_message.metrics.start_timer()

            # Track current tool being built across chunks
            current_tool: Dict[str, Any] = {}

            async with self.get_async_client() as client:
                response = await client.converse_stream(modelId=self.id, messages=formatted_messages, **body)
                async for chunk in response["stream"]:
                    model_response, current_tool = self._parse_provider_response_delta(chunk, current_tool)
                    yield model_response

            assistant_message.metrics.stop_timer()

        except ClientError as e:
            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
            raise ModelProviderError(message=str(e.response), model_name=self.name, model_id=self.id) from e
        except Exception as e:
            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
            raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e

    # Overwrite the default from the base model
    def format_function_call_results(
        self,
        messages: List[Message],
        function_call_results: List[Message],
        compress_tool_results: bool = False,
        **kwargs,
    ) -> None:
        """
        Handle the results of function calls for Bedrock.
        Uses compressed_content if compress_tool_results is True.

        Args:
            messages (List[Message]): The list of conversation messages.
            function_call_results (List[Message]): The results of the function calls.
            compress_tool_results: Whether to compress tool results.
            **kwargs: Additional arguments including tool_ids.
        """

        if function_call_results:
            tool_ids = kwargs.get("tool_ids", [])

            for _fc_message_index, _fc_message in enumerate(function_call_results):
                # Use tool_call_id from message if tool_ids list is insufficient
                tool_id = tool_ids[_fc_message_index] if _fc_message_index < len(tool_ids) else _fc_message.tool_call_id
                if not _fc_message.tool_call_id:
                    _fc_message.tool_call_id = tool_id

                # Append as standard role="tool" message
                messages.append(_fc_message)

    def _parse_provider_response(self, response: Dict[str, Any], **kwargs) -> ModelResponse:
        """
        Parse the provider response.

        Args:
            response (Dict[str, Any]): The response from the provider.

        Returns:
            ModelResponse: The parsed response.
        """
        model_response = ModelResponse()

        if "output" in response and "message" in response["output"]:
            message = response["output"]["message"]
            # Set the role of the message
            model_response.role = message["role"]

            # Get the content of the message
            content = message["content"]

            # Tools
            if "stopReason" in response and response["stopReason"] == "tool_use":
                model_response.tool_calls = []
                model_response.extra = model_response.extra or {}
                model_response.extra["tool_ids"] = []
                for tool in content:
                    if "toolUse" in tool:
                        model_response.extra["tool_ids"].append(tool["toolUse"]["toolUseId"])
                        model_response.tool_calls.append(
                            {
                                "id": tool["toolUse"]["toolUseId"],
                                "type": "function",
                                "function": {
                                    "name": tool["toolUse"]["name"],
                                    "arguments": json.dumps(tool["toolUse"]["input"]),
                                },
                            }
                        )

            # Extract text content if it's a list of dictionaries
            if isinstance(content, list) and content and isinstance(content[0], dict):
                content = [item.get("text", "") for item in content if "text" in item]
                content = "\n".join(content)  # Join multiple text items if present

            model_response.content = content

        if "usage" in response:
            model_response.response_usage = self._get_metrics(response["usage"])

        return model_response

    def _parse_provider_response_delta(
        self, response_delta: Dict[str, Any], current_tool: Dict[str, Any]
    ) -> Tuple[ModelResponse, Dict[str, Any]]:
        """Parse the provider response delta for streaming.

        Args:
            response_delta: The streaming response delta from AWS Bedrock
            current_tool: The current tool being built across chunks

        Returns:
            Tuple[ModelResponse, Dict[str, Any]]: The parsed model response delta and updated current_tool
        """
        model_response = ModelResponse(role="assistant")

        # Handle contentBlockStart - tool use start
        if "contentBlockStart" in response_delta:
            start = response_delta["contentBlockStart"]["start"]
            if "toolUse" in start:
                # Start a new tool
                tool_use_data = start["toolUse"]
                current_tool = {
                    "id": tool_use_data.get("toolUseId", ""),
                    "type": "function",
                    "function": {
                        "name": tool_use_data.get("name", ""),
                        "arguments": "",  # Will be filled in subsequent deltas
                    },
                }

        # Handle contentBlockDelta - text content or tool input
        elif "contentBlockDelta" in response_delta:
            delta = response_delta["contentBlockDelta"]["delta"]
            if "text" in delta:
                model_response.content = delta["text"]
            elif "toolUse" in delta and current_tool:
                # Accumulate tool input
                tool_input = delta["toolUse"].get("input", "")
                if tool_input:
                    current_tool["function"]["arguments"] += tool_input

        # Handle contentBlockStop - tool use complete
        elif "contentBlockStop" in response_delta and current_tool:
            # Tool is complete, add it to model response
            model_response.tool_calls = [current_tool]
            # Track tool_id in extra for format_function_call_results
            model_response.extra = {"tool_ids": [current_tool["id"]]}
            # Reset current_tool for next tool
            current_tool = {}

        # Handle metadata/usage information
        elif "metadata" in response_delta or "messageStop" in response_delta:
            body = response_delta.get("metadata") or response_delta.get("messageStop") or {}
            if "usage" in body:
                model_response.response_usage = self._get_metrics(body["usage"])

        return model_response, current_tool

    def _get_metrics(self, response_usage: Dict[str, Any]) -> MessageMetrics:
        """
        Parse the given AWS Bedrock usage into an Agno MessageMetrics object.

        Args:
            response_usage: Usage data from AWS Bedrock

        Returns:
            MessageMetrics: Parsed metrics data
        """
        metrics = MessageMetrics()

        metrics.input_tokens = response_usage.get("inputTokens", 0) or 0
        metrics.output_tokens = response_usage.get("outputTokens", 0) or 0
        metrics.total_tokens = metrics.input_tokens + metrics.output_tokens

        metrics.cache_read_tokens = response_usage.get("cacheReadInputTokens", 0) or 0
        metrics.cache_write_tokens = response_usage.get("cacheWriteInputTokens", 0) or 0

        return metrics
