# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0
#
# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
# SPDX-License-Identifier: MIT
"""Create a compatible client for the Amazon Bedrock Converse API.

Example usage:
Install the `boto3` package by running `pip install --upgrade boto3`.
- https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html

```python
import autogen

config_list = [
    {
        "api_type": "bedrock",
        "model": "meta.llama3-1-8b-instruct-v1:0",
        "aws_region": "us-west-2",
        "aws_access_key": "",
        "aws_secret_key": "",
        "price": [0.003, 0.015],
    }
]

assistant = autogen.AssistantAgent("assistant", llm_config={"config_list": config_list})
```
"""

from __future__ import annotations

import base64
import json
import os
import re
import time
import warnings
from typing import Any, Literal, Optional

import requests
from pydantic import Field, SecretStr, field_serializer

from ..import_utils import optional_import_block, require_optional_import
from ..llm_config import LLMConfigEntry, register_llm_config
from .client_utils import validate_parameter
from .oai_models import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageToolCall, Choice, CompletionUsage

with optional_import_block():
    import boto3
    from botocore.config import Config


@register_llm_config
class BedrockLLMConfigEntry(LLMConfigEntry):
    api_type: Literal["bedrock"] = "bedrock"
    aws_region: str
    aws_access_key: Optional[SecretStr] = None
    aws_secret_key: Optional[SecretStr] = None
    aws_session_token: Optional[SecretStr] = None
    aws_profile_name: Optional[str] = None
    temperature: Optional[float] = None
    topP: Optional[float] = None  # noqa: N815
    maxTokens: Optional[int] = None  # noqa: N815
    top_p: Optional[float] = None
    top_k: Optional[int] = None
    k: Optional[int] = None
    seed: Optional[int] = None
    cache_seed: Optional[int] = None
    supports_system_prompts: bool = True
    stream: bool = False
    price: Optional[list[float]] = Field(default=None, min_length=2, max_length=2)
    timeout: Optional[int] = None

    @field_serializer("aws_access_key", "aws_secret_key", "aws_session_token", when_used="unless-none")
    def serialize_aws_secrets(self, v: SecretStr) -> str:
        return v.get_secret_value()

    def create_client(self):
        raise NotImplementedError("BedrockLLMConfigEntry.create_client must be implemented.")


@require_optional_import("boto3", "bedrock")
class BedrockClient:
    """Client for Amazon's Bedrock Converse API."""

    _retries = 5

    def __init__(self, **kwargs: Any):
        """Initialises BedrockClient for Amazon's Bedrock Converse API"""
        self._aws_access_key = kwargs.get("aws_access_key")
        self._aws_secret_key = kwargs.get("aws_secret_key")
        self._aws_session_token = kwargs.get("aws_session_token")
        self._aws_region = kwargs.get("aws_region")
        self._aws_profile_name = kwargs.get("aws_profile_name")
        self._timeout = kwargs.get("timeout")

        if not self._aws_access_key:
            self._aws_access_key = os.getenv("AWS_ACCESS_KEY")

        if not self._aws_secret_key:
            self._aws_secret_key = os.getenv("AWS_SECRET_KEY")

        if not self._aws_session_token:
            self._aws_session_token = os.getenv("AWS_SESSION_TOKEN")

        if not self._aws_region:
            self._aws_region = os.getenv("AWS_REGION")

        if self._aws_region is None:
            raise ValueError("Region is required to use the Amazon Bedrock API.")

        if self._timeout is None:
            self._timeout = 60

        # Initialize Bedrock client, session, and runtime
        bedrock_config = Config(
            region_name=self._aws_region,
            signature_version="v4",
            retries={"max_attempts": self._retries, "mode": "standard"},
            read_timeout=self._timeout,
        )

        session = boto3.Session(
            aws_access_key_id=self._aws_access_key,
            aws_secret_access_key=self._aws_secret_key,
            aws_session_token=self._aws_session_token,
            profile_name=self._aws_profile_name,
        )

        if "response_format" in kwargs and kwargs["response_format"] is not None:
            warnings.warn("response_format is not supported for Bedrock, it will be ignored.", UserWarning)

        # if haven't got any access_key or secret_key in environment variable or via arguments then
        if (
            self._aws_access_key is None
            or self._aws_access_key == ""
            or self._aws_secret_key is None
            or self._aws_secret_key == ""
        ):
            # attempts to get client from attached role of managed service (lambda, ec2, ecs, etc.)
            self.bedrock_runtime = boto3.client(service_name="bedrock-runtime", config=bedrock_config)
        else:
            session = boto3.Session(
                aws_access_key_id=self._aws_access_key,
                aws_secret_access_key=self._aws_secret_key,
                aws_session_token=self._aws_session_token,
                profile_name=self._aws_profile_name,
            )
            self.bedrock_runtime = session.client(service_name="bedrock-runtime", config=bedrock_config)

    def message_retrieval(self, response):
        """Retrieve the messages from the response."""
        return [choice.message for choice in response.choices]

    def parse_custom_params(self, params: dict[str, Any]):
        """Parses custom parameters for logic in this client class"""
        # Should we separate system messages into its own request parameter, default is True
        # This is required because not all models support a system prompt (e.g. Mistral Instruct).
        self._supports_system_prompts = params.get("supports_system_prompts", True)

    def parse_params(self, params: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
        """Loads the valid parameters required to invoke Bedrock Converse
        Returns a tuple of (base_params, additional_params)
        """
        base_params = {}
        additional_params = {}

        # Amazon Bedrock  base model IDs are here:
        # https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html
        self._model_id = params.get("model")
        assert self._model_id, "Please provide the 'model` in the config_list to use Amazon Bedrock"

        # Parameters vary based on the model used.
        # As we won't cater for all models and parameters, it's the developer's
        # responsibility to implement the parameters and they will only be
        # included if the developer has it in the config.
        #
        # Important:
        # No defaults will be used (as they can vary per model)
        # No ranges will be used (as they can vary)
        # We will cover all the main parameters but there may be others
        # that need to be added later
        #
        # Here are some pages that show the parameters available for different models
        # https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-text.html
        # https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-text-completion.html
        # https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-command-r-plus.html
        # https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-meta.html
        # https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mistral-chat-completion.html

        # Here are the possible "base" parameters and their suitable types
        base_parameters = [["temperature", (float, int)], ["topP", (float, int)], ["maxTokens", (int)]]

        for param_name, suitable_types in base_parameters:
            if param_name in params:
                base_params[param_name] = validate_parameter(
                    params, param_name, suitable_types, False, None, None, None
                )

        # Here are the possible "model-specific" parameters and their suitable types, known as additional parameters
        additional_parameters = [
            ["top_p", (float, int)],
            ["top_k", (int)],
            ["k", (int)],
            ["seed", (int)],
        ]

        for param_name, suitable_types in additional_parameters:
            if param_name in params:
                additional_params[param_name] = validate_parameter(
                    params, param_name, suitable_types, False, None, None, None
                )

        # Streaming
        self._streaming = params.get("stream", False)

        # For this release we will not support streaming as many models do not support streaming with tool use
        if self._streaming:
            warnings.warn(
                "Streaming is not currently supported, streaming will be disabled.",
                UserWarning,
            )
            self._streaming = False

        return base_params, additional_params

    def create(self, params) -> ChatCompletion:
        """Run Amazon Bedrock inference and return AG2 response"""
        # Set custom client class settings
        self.parse_custom_params(params)

        # Parse the inference parameters
        base_params, additional_params = self.parse_params(params)

        has_tools = "tools" in params
        messages = oai_messages_to_bedrock_messages(params["messages"], has_tools, self._supports_system_prompts)

        if self._supports_system_prompts:
            system_messages = extract_system_messages(params["messages"])

        tool_config = format_tools(params["tools"] if has_tools else [])

        request_args = {"messages": messages, "modelId": self._model_id}

        # Base and additional args
        if len(base_params) > 0:
            request_args["inferenceConfig"] = base_params

        if len(additional_params) > 0:
            request_args["additionalModelRequestFields"] = additional_params

        if self._supports_system_prompts:
            request_args["system"] = system_messages

        if len(tool_config["tools"]) > 0:
            request_args["toolConfig"] = tool_config

        response = self.bedrock_runtime.converse(**request_args)
        if response is None:
            raise RuntimeError(f"Failed to get response from Bedrock after retrying {self._retries} times.")

        finish_reason = convert_stop_reason_to_finish_reason(response["stopReason"])
        response_message = response["output"]["message"]

        tool_calls = format_tool_calls(response_message["content"]) if finish_reason == "tool_calls" else None

        text = ""
        for content in response_message["content"]:
            if "text" in content:
                text = content["text"]
                # NOTE: other types of output may be dealt with here

        message = ChatCompletionMessage(role="assistant", content=text, tool_calls=tool_calls)

        response_usage = response["usage"]
        usage = CompletionUsage(
            prompt_tokens=response_usage["inputTokens"],
            completion_tokens=response_usage["outputTokens"],
            total_tokens=response_usage["totalTokens"],
        )

        return ChatCompletion(
            id=response["ResponseMetadata"]["RequestId"],
            choices=[Choice(finish_reason=finish_reason, index=0, message=message)],
            created=int(time.time()),
            model=self._model_id,
            object="chat.completion",
            usage=usage,
        )

    def cost(self, response: ChatCompletion) -> float:
        """Calculate the cost of the response."""
        return calculate_cost(response.usage.prompt_tokens, response.usage.completion_tokens, response.model)

    @staticmethod
    def get_usage(response) -> dict:
        """Get the usage of tokens and their cost information."""
        return {
            "prompt_tokens": response.usage.prompt_tokens,
            "completion_tokens": response.usage.completion_tokens,
            "total_tokens": response.usage.total_tokens,
            "cost": response.cost,
            "model": response.model,
        }


def extract_system_messages(messages: list[dict[str, Any]]) -> list:
    """Extract the system messages from the list of messages.

    Args:
        messages (list[dict[str, Any]]): List of messages.

    Returns:
        List[SystemMessage]: List of System messages.
    """
    """
    system_messages = [message.get("content")[0]["text"] for message in messages if message.get("role") == "system"]
    return system_messages # ''.join(system_messages)
    """

    for message in messages:
        if message.get("role") == "system":
            if isinstance(message["content"], str):
                return [{"text": message.get("content")}]
            else:
                return [{"text": message.get("content")[0]["text"]}]
    return []


def oai_messages_to_bedrock_messages(
    messages: list[dict[str, Any]], has_tools: bool, supports_system_prompts: bool
) -> list[dict[str, Any]]:
    """Convert messages from OAI format to Bedrock format.
    We correct for any specific role orders and types, etc.
    AWS Bedrock requires messages to alternate between user and assistant roles. This function ensures that the messages
    are in the correct order and format for Bedrock by inserting "Please continue" messages as needed.
    This is the same method as the one in the Autogen Anthropic client
    """
    # Track whether we have tools passed in. If not,  tool use / result messages should be converted to text messages.
    # Bedrock requires a tools parameter with the tools listed, if there are other messages with tool use or tool results.
    # This can occur when we don't need tool calling, such as for group chat speaker selection

    # Convert messages to Bedrock compliant format

    # Take out system messages if the model supports it, otherwise leave them in.
    if supports_system_prompts:
        messages = [x for x in messages if x["role"] != "system"]
    else:
        # Replace role="system" with role="user"
        for msg in messages:
            if msg["role"] == "system":
                msg["role"] = "user"

    processed_messages = []

    # Used to interweave user messages to ensure user/assistant alternating
    user_continue_message = {"content": [{"text": "Please continue."}], "role": "user"}
    assistant_continue_message = {
        "content": [{"text": "Please continue."}],
        "role": "assistant",
    }

    tool_use_messages = 0
    tool_result_messages = 0
    last_tool_use_index = -1
    last_tool_result_index = -1
    # user_role_index = 0 if supports_system_prompts else 1 # If system prompts are supported, messages start with user, otherwise they'll be the second message
    for message in messages:
        # New messages will be added here, manage role alternations
        expected_role = "user" if len(processed_messages) % 2 == 0 else "assistant"

        if "tool_calls" in message:
            # Map the tool call options to Bedrock's format
            tool_uses = []
            tool_names = []
            for tool_call in message["tool_calls"]:
                tool_uses.append({
                    "toolUse": {
                        "toolUseId": tool_call["id"],
                        "name": tool_call["function"]["name"],
                        "input": json.loads(tool_call["function"]["arguments"]),
                    }
                })
                if has_tools:
                    tool_use_messages += 1
                tool_names.append(tool_call["function"]["name"])

            if expected_role == "user":
                # Insert an extra user message as we will append an assistant message
                processed_messages.append(user_continue_message)

            if has_tools:
                processed_messages.append({"role": "assistant", "content": tool_uses})
                last_tool_use_index = len(processed_messages) - 1
            else:
                # Not using tools, so put in a plain text message
                processed_messages.append({
                    "role": "assistant",
                    "content": [{"text": f"Some internal function(s) that could be used: [{', '.join(tool_names)}]"}],
                })
        elif "tool_call_id" in message:
            if has_tools:
                # Map the tool usage call to tool_result for Bedrock
                tool_result = {
                    "toolResult": {
                        "toolUseId": message["tool_call_id"],
                        "content": [{"text": message["content"]}],
                    }
                }

                # If the previous message also had a tool_result, add it to that
                # Otherwise append a new message
                if last_tool_result_index == len(processed_messages) - 1:
                    processed_messages[-1]["content"].append(tool_result)
                else:
                    if expected_role == "assistant":
                        # Insert an extra assistant message as we will append a user message
                        processed_messages.append(assistant_continue_message)

                    processed_messages.append({"role": "user", "content": [tool_result]})
                    last_tool_result_index = len(processed_messages) - 1

                tool_result_messages += 1
            else:
                # Not using tools, so put in a plain text message
                processed_messages.append({
                    "role": "user",
                    "content": [{"text": f"Running the function returned: {message['content']}"}],
                })
        elif message["content"] == "":
            # Ignoring empty messages
            pass
        else:
            if expected_role != message["role"] and not (len(processed_messages) == 0 and message["role"] == "system"):
                # Inserting the alternating continue message (ignore if it's the first message and a system message)
                processed_messages.append(
                    user_continue_message if expected_role == "user" else assistant_continue_message
                )

            processed_messages.append({
                "role": message["role"],
                "content": parse_content_parts(message=message),
            })

    # We'll replace the last tool_use if there's no tool_result (occurs if we finish the conversation before running the function)
    if has_tools and tool_use_messages != tool_result_messages:
        processed_messages[last_tool_use_index] = assistant_continue_message

    # name is not a valid field on messages
    for message in processed_messages:
        if "name" in message:
            message.pop("name", None)

    # Note: When using reflection_with_llm we may end up with an "assistant" message as the last message and that may cause a blank response
    # So, if the last role is not user, add a 'user' continue message at the end
    if processed_messages[-1]["role"] != "user":
        processed_messages.append(user_continue_message)

    return processed_messages


def parse_content_parts(
    message: dict[str, Any],
) -> list[dict[str, Any]]:
    content: str | list[dict[str, Any]] = message.get("content")
    if isinstance(content, str):
        return [
            {
                "text": content,
            }
        ]
    content_parts = []
    for part in content:
        # part_content: Dict = part.get("content")
        if "text" in part:  # part_content:
            content_parts.append({
                "text": part.get("text"),
            })
        elif "image_url" in part:  # part_content:
            image_data, content_type = parse_image(part.get("image_url").get("url"))
            content_parts.append({
                "image": {
                    "format": content_type[6:],  # image/
                    "source": {"bytes": image_data},
                },
            })
        else:
            # Ignore..
            continue
    return content_parts


def parse_image(image_url: str) -> tuple[bytes, str]:
    """Try to get the raw data from an image url.

    Ref: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ImageSource.html
    returns a tuple of (Image Data, Content Type)
    """
    pattern = r"^data:(image/[a-z]*);base64,\s*"
    content_type = re.search(pattern, image_url)
    # if already base64 encoded.
    # Only supports 'image/jpeg', 'image/png', 'image/gif' or 'image/webp'
    if content_type:
        image_data = re.sub(pattern, "", image_url)
        return base64.b64decode(image_data), content_type.group(1)

    # Send a request to the image URL
    response = requests.get(image_url)
    # Check if the request was successful
    if response.status_code == 200:
        content_type = response.headers.get("Content-Type")
        if not content_type.startswith("image"):
            content_type = "image/jpeg"
        # Get the image content
        image_content = response.content
        return image_content, content_type
    else:
        raise RuntimeError("Unable to access the image url")


def format_tools(tools: list[dict[str, Any]]) -> dict[Literal["tools"], list[dict[str, Any]]]:
    converted_schema = {"tools": []}

    for tool in tools:
        if tool["type"] == "function":
            function = tool["function"]
            converted_tool = {
                "toolSpec": {
                    "name": function["name"],
                    "description": function["description"],
                    "inputSchema": {"json": {"type": "object", "properties": {}, "required": []}},
                }
            }

            for prop_name, prop_details in function["parameters"]["properties"].items():
                converted_tool["toolSpec"]["inputSchema"]["json"]["properties"][prop_name] = {
                    "type": prop_details["type"],
                    "description": prop_details.get("description", ""),
                }
                if "enum" in prop_details:
                    converted_tool["toolSpec"]["inputSchema"]["json"]["properties"][prop_name]["enum"] = prop_details[
                        "enum"
                    ]
                if "default" in prop_details:
                    converted_tool["toolSpec"]["inputSchema"]["json"]["properties"][prop_name]["default"] = (
                        prop_details["default"]
                    )

            if "required" in function["parameters"]:
                converted_tool["toolSpec"]["inputSchema"]["json"]["required"] = function["parameters"]["required"]

            converted_schema["tools"].append(converted_tool)

    return converted_schema


def format_tool_calls(content):
    """Converts Converse API response tool calls to AG2 format"""
    tool_calls = []
    for tool_request in content:
        if "toolUse" in tool_request:
            tool = tool_request["toolUse"]

            tool_calls.append(
                ChatCompletionMessageToolCall(
                    id=tool["toolUseId"],
                    function={
                        "name": tool["name"],
                        "arguments": json.dumps(tool["input"]),
                    },
                    type="function",
                )
            )
    return tool_calls


def convert_stop_reason_to_finish_reason(
    stop_reason: str,
) -> Literal["stop", "length", "tool_calls", "content_filter"]:
    """Converts Bedrock finish reasons to our finish reasons, according to OpenAI:

    - stop: if the model hit a natural stop point or a provided stop sequence,
    - length: if the maximum number of tokens specified in the request was reached,
    - content_filter: if content was omitted due to a flag from our content filters,
    - tool_calls: if the model called a tool
    """
    if stop_reason:
        finish_reason_mapping = {
            "tool_use": "tool_calls",
            "finished": "stop",
            "end_turn": "stop",
            "max_tokens": "length",
            "stop_sequence": "stop",
            "complete": "stop",
            "content_filtered": "content_filter",
        }
        return finish_reason_mapping.get(stop_reason.lower(), stop_reason.lower())

    warnings.warn(f"Unsupported stop reason: {stop_reason}", UserWarning)
    return None


# NOTE: As this will be quite dynamic, it's expected that the developer will use the "price" parameter in their config
# These may be removed.
PRICES_PER_K_TOKENS = {
    "meta.llama3-8b-instruct-v1:0": (0.0003, 0.0006),
    "meta.llama3-70b-instruct-v1:0": (0.00265, 0.0035),
    "mistral.mistral-7b-instruct-v0:2": (0.00015, 0.0002),
    "mistral.mixtral-8x7b-instruct-v0:1": (0.00045, 0.0007),
    "mistral.mistral-large-2402-v1:0": (0.004, 0.012),
    "mistral.mistral-small-2402-v1:0": (0.001, 0.003),
}


def calculate_cost(input_tokens: int, output_tokens: int, model_id: str) -> float:
    """Calculate the cost of the completion using the Bedrock pricing."""
    if model_id in PRICES_PER_K_TOKENS:
        input_cost_per_k, output_cost_per_k = PRICES_PER_K_TOKENS[model_id]
        input_cost = (input_tokens / 1000) * input_cost_per_k
        output_cost = (output_tokens / 1000) * output_cost_per_k
        return input_cost + output_cost
    else:
        warnings.warn(
            f'Cannot get the costs for {model_id}. The cost will be 0. In your config_list, add field {{"price" : [prompt_price_per_1k, completion_token_price_per_1k]}} for customized pricing.',
            UserWarning,
        )
        return 0