import os

from autogen import ConversableAgent


class RELAIAgent(ConversableAgent):
    """Base class for interactive RELAI agents that can participate in User-LLM converations.
    All RELAI agents must inherit from this class.

    Args:
        model (str): Model name to use for the agent. Default is "gpt-4o".
        llm_config (dict): LLM config to use for the agent. If provided, `model` argument is ignored. If not provided, it will be fetched based on the model.
        *args: Additional arguments to pass to the parent class.
        **kwargs: Additional keyword arguments to pass to the parent class.
    """

    def __init__(self, model: str = "gpt-4o", llm_config=None, *args, **kwargs) -> None:
        llm_config = self._get_llm_config(model) if not llm_config else llm_config
        kwargs["llm_config"] = {
            "config_list": [llm_config],
        }
        if "name" not in kwargs:
            # Set default name for RELAI agents
            kwargs["name"] = "relai-agent"
        super().__init__(*args, **kwargs)

    def _get_llm_config(self, model: str):
        """Util to get `llm_config` for a given model.
        Note that it is a private method since config has secrets values.
        """
        aws_access_key = os.environ.get("AWS_ACCESS_KEY_ID", None)
        aws_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY", None)
        gemini_api_key = os.environ.get("GEMINI_API_KEY", None)

        if not aws_access_key or not aws_secret_key:
            raise ValueError("Please set `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` env vars.")

        if (
            "gpt-4o" in model or "gpt-4o-mini" in model or "o1" in model or "o1-mini" in model or "o3-mini" in model
        ):  # models fine-tuned on OpenAI have base model in their name
            return {
                "model": model,
            }
        elif model == "bedrock/meta.llama3-1-8b-instruct-v1:0":
            return {
                "api_type": "bedrock",
                "model": "meta.llama3-1-8b-instruct-v1:0",
                "aws_region": "us-west-2",
                "aws_access_key": aws_access_key,
                "aws_secret_key": aws_secret_key,
                "price": [0.00022, 0.00022],
            }
        elif model == "bedrock/meta.llama3-1-70b-instruct-v1:0":
            return {
                "api_type": "bedrock",
                "model": "meta.llama3-1-70b-instruct-v1:0",
                "aws_region": "us-west-2",
                "aws_access_key": aws_access_key,
                "aws_secret_key": aws_secret_key,
                "price": [0.00099, 0.00099],
            }
        elif model == "bedrock/meta.llama3-1-405b-instruct-v1:0":
            return {
                "api_type": "bedrock",
                "model": "meta.llama3-1-405b-instruct-v1:0",
                "aws_region": "us-west-2",
                "aws_access_key": aws_access_key,
                "aws_secret_key": aws_secret_key,
                "price": [0.00532, 0.016],
            }
        elif model == "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0":
            return {
                "api_type": "bedrock",
                "model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
                "aws_region": "us-east-1",
                "aws_access_key": aws_access_key,
                "aws_secret_key": aws_secret_key,
                "price": [0.003, 0.015],
            }
        elif model == "bedrock/meta.llama3-3-70b-instruct-v1:0":
            return {
                "api_type": "bedrock",
                "model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
                "aws_region": "us-east-1",
                "aws_access_key": aws_access_key,
                "aws_secret_key": aws_secret_key,
                "price": [0.00072, 0.00072],
            }
        elif model == "gemini/gemini-pro":
            return {
                "api_type": "google",
                "model": "gemini-pro",
                "api_key": gemini_api_key,
                "price": [0.00125, 0.005],
            }
        else:
            raise ValueError(f"`{model}` is not supported.")

    def construct_agent_input(self, messages: list[dict]):
        """Combines human prompts and AI responses into a format expected by an agent."""
        history_len = 2  # Number of previous messages to include in history
        messages = messages[
            -2 * (history_len + 1) :
        ]  # Limit to last `history_len` messages and current prompt/response

        history = [{} for _ in range(history_len)]
        for idx, message in enumerate(messages[:-2]):
            if message["role"] == "user":
                history[idx // 2]["prompt"] = message["content"]
            elif message["role"] == "assistant":
                history[idx // 2]["response"] = message["content"]

        history = history[::-1]  # Reverse order of history so that most recent message is first

        content = f"[PROMPT]:\n{messages[-2]['content']}\n\n[RESPONSE]:\n{messages[-1]['content']}\n\n"
        content += f"[HISTORY]:\n{history}"
        return [{"role": "user", "content": content}]
