import json
import os
import re
from typing import Any, Dict, List, Optional

try:
    import boto3
except ImportError:
    raise ImportError("The 'boto3' library is required. Please install it using 'pip install boto3'.")

from mem0.configs.llms.base import BaseLlmConfig
from mem0.llms.base import LLMBase

PROVIDERS = ["ai21", "amazon", "anthropic", "cohere", "meta", "mistral", "stability", "writer"]


def extract_provider(model: str) -> str:
    for provider in PROVIDERS:
        if re.search(rf"\b{re.escape(provider)}\b", model):
            return provider
    raise ValueError(f"Unknown provider in model: {model}")


class AWSBedrockLLM(LLMBase):
    def __init__(self, config: Optional[BaseLlmConfig] = None):
        super().__init__(config)

        if not self.config.model:
            self.config.model = "anthropic.claude-3-5-sonnet-20240620-v1:0"

        # Get AWS config from environment variables or use defaults
        aws_access_key = os.environ.get("AWS_ACCESS_KEY_ID", "")
        aws_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY", "")
        aws_region = os.environ.get("AWS_REGION", "us-west-2")

        # Check if AWS config is provided in the config
        if hasattr(self.config, "aws_access_key_id"):
            aws_access_key = self.config.aws_access_key_id
        if hasattr(self.config, "aws_secret_access_key"):
            aws_secret_key = self.config.aws_secret_access_key
        if hasattr(self.config, "aws_region"):
            aws_region = self.config.aws_region

        self.client = boto3.client(
            "bedrock-runtime",
            region_name=aws_region,
            aws_access_key_id=aws_access_key if aws_access_key else None,
            aws_secret_access_key=aws_secret_key if aws_secret_key else None,
        )

        self.model_kwargs = {
            "temperature": self.config.temperature,
            "max_tokens_to_sample": self.config.max_tokens,
            "top_p": self.config.top_p,
        }

    def _format_messages(self, messages: List[Dict[str, str]]) -> str:
        """
        Formats a list of messages into the required prompt structure for the model.

        Args:
            messages (List[Dict[str, str]]): A list of dictionaries where each dictionary represents a message.
                                            Each dictionary contains 'role' and 'content' keys.

        Returns:
            str: A formatted string combining all messages, structured with roles capitalized and separated by newlines.
        """

        formatted_messages = []
        for message in messages:
            role = message["role"].capitalize()
            content = message["content"]
            formatted_messages.append(f"\n\n{role}: {content}")

        return "\n\nHuman: " + "".join(formatted_messages) + "\n\nAssistant:"

    def _parse_response(self, response, tools) -> str:
        """
        Process the response based on whether tools are used or not.

        Args:
            response: The raw response from API.
            tools: The list of tools provided in the request.

        Returns:
            str or dict: The processed response.
        """
        if tools:
            processed_response = {"tool_calls": []}

            if response["output"]["message"]["content"]:
                for item in response["output"]["message"]["content"]:
                    if "toolUse" in item:
                        processed_response["tool_calls"].append(
                            {
                                "name": item["toolUse"]["name"],
                                "arguments": item["toolUse"]["input"],
                            }
                        )

            return processed_response

        response_body = response.get("body").read().decode()
        response_json = json.loads(response_body)
        return response_json.get("content", [{"text": ""}])[0].get("text", "")

    def _prepare_input(
        self,
        provider: str,
        model: str,
        prompt: str,
        model_kwargs: Optional[Dict[str, Any]] = {},
    ) -> Dict[str, Any]:
        """
        Prepares the input dictionary for the specified provider's model by mapping and renaming
        keys in the input based on the provider's requirements.

        Args:
            provider (str): The name of the service provider (e.g., "meta", "ai21", "mistral", "cohere", "amazon").
            model (str): The name or identifier of the model being used.
            prompt (str): The text prompt to be processed by the model.
            model_kwargs (Dict[str, Any]): Additional keyword arguments specific to the model's requirements.

        Returns:
            Dict[str, Any]: The prepared input dictionary with the correct keys and values for the specified provider.
        """

        input_body = {"prompt": prompt, **model_kwargs}

        provider_mappings = {
            "meta": {"max_tokens_to_sample": "max_gen_len"},
            "ai21": {"max_tokens_to_sample": "maxTokens", "top_p": "topP"},
            "mistral": {"max_tokens_to_sample": "max_tokens"},
            "cohere": {"max_tokens_to_sample": "max_tokens", "top_p": "p"},
        }

        if provider in provider_mappings:
            for old_key, new_key in provider_mappings[provider].items():
                if old_key in input_body:
                    input_body[new_key] = input_body.pop(old_key)

        if provider == "cohere" and "cohere.command-r" in model:
            input_body["message"] = input_body.pop("prompt")

        if provider == "amazon":
            input_body = {
                "inputText": prompt,
                "textGenerationConfig": {
                    "maxTokenCount": self.model_kwargs["max_tokens_to_sample"]
                    or self.model_kwargs["max_tokens"]
                    or 5000,
                    "topP": self.model_kwargs["top_p"] or 0.9,
                    "temperature": self.model_kwargs["temperature"] or 0.1,
                },
            }
            input_body["textGenerationConfig"] = {
                k: v for k, v in input_body["textGenerationConfig"].items() if v is not None
            }

        return input_body

    def _convert_tool_format(self, original_tools):
        """
        Converts a list of tools from their original format to a new standardized format.

        Args:
            original_tools (list): A list of dictionaries representing the original tools, each containing a 'type' key and corresponding details.

        Returns:
            list: A list of dictionaries representing the tools in the new standardized format.
        """
        new_tools = []

        for tool in original_tools:
            if tool["type"] == "function":
                function = tool["function"]
                new_tool = {
                    "toolSpec": {
                        "name": function["name"],
                        "description": function["description"],
                        "inputSchema": {
                            "json": {
                                "type": "object",
                                "properties": {},
                                "required": function["parameters"].get("required", []),
                            }
                        },
                    }
                }

                for prop, details in function["parameters"].get("properties", {}).items():
                    new_tool["toolSpec"]["inputSchema"]["json"]["properties"][prop] = {
                        "type": details.get("type", "string"),
                        "description": details.get("description", ""),
                    }

                new_tools.append(new_tool)

        return new_tools

    def generate_response(
        self,
        messages: List[Dict[str, str]],
        response_format=None,
        tools: Optional[List[Dict]] = None,
        tool_choice: str = "auto",
    ):
        """
        Generate a response based on the given messages using AWS Bedrock.

        Args:
            messages (list): List of message dicts containing 'role' and 'content'.
            tools (list, optional): List of tools that the model can call. Defaults to None.
            tool_choice (str, optional): Tool choice method. Defaults to "auto".

        Returns:
            str: The generated response.
        """

        if tools:
            # Use converse method when tools are provided
            messages = [
                {
                    "role": "user",
                    "content": [{"text": message["content"]} for message in messages],
                }
            ]
            inference_config = {
                "temperature": self.model_kwargs["temperature"],
                "maxTokens": self.model_kwargs["max_tokens_to_sample"],
                "topP": self.model_kwargs["top_p"],
            }
            tools_config = {"tools": self._convert_tool_format(tools)}

            response = self.client.converse(
                modelId=self.config.model,
                messages=messages,
                inferenceConfig=inference_config,
                toolConfig=tools_config,
            )
        else:
            # Use invoke_model method when no tools are provided
            prompt = self._format_messages(messages)
            provider = extract_provider(self.config.model)
            input_body = self._prepare_input(provider, self.config.model, prompt, model_kwargs=self.model_kwargs)
            body = json.dumps(input_body)

            if provider == "anthropic" or provider == "deepseek":
                input_body = {
                    "messages": [{"role": "user", "content": [{"type": "text", "text": prompt}]}],
                    "max_tokens": self.model_kwargs["max_tokens_to_sample"] or self.model_kwargs["max_tokens"] or 5000,
                    "temperature": self.model_kwargs["temperature"] or 0.1,
                    "top_p": self.model_kwargs["top_p"] or 0.9,
                    "anthropic_version": "bedrock-2023-05-31",
                }

                body = json.dumps(input_body)

                response = self.client.invoke_model(
                    body=body,
                    modelId=self.config.model,
                    accept="application/json",
                    contentType="application/json",
                )
            else:
                response = self.client.invoke_model(
                    body=body,
                    modelId=self.config.model,
                    accept="application/json",
                    contentType="application/json",
                )

        return self._parse_response(response, tools)
