import asyncio
import json
import logging
from dataclasses import dataclass, field
from functools import partial
from typing import Any, Dict, List, Optional, Union

from phoenix.evals.exceptions import PhoenixContextLimitExceeded
from phoenix.evals.models.base import BaseModel
from phoenix.evals.models.rate_limiters import RateLimiter
from phoenix.evals.templates import MultimodalPrompt

logger = logging.getLogger(__name__)

MINIMUM_BOTO_VERSION = "1.28.58"


@dataclass
class BedrockModel(BaseModel):
    """
    An interface for using LLM models via AWS Bedrock.

    This class wraps the boto3 Bedrock client for use with Phoenix LLM evaluations. Calls to the
    AWS API are dynamically throttled when encountering rate limit errors. Requires the `boto3`
    package to be installed.

    Supports Async: 🟡
        `boto3` does not support async calls, so it's wrapped in an executor.

    Args:
        model_id (str): The model name to use.
        temperature (float, optional): Sampling temperature to use. Defaults to 0.0.
        max_tokens (int, optional): Maximum number of tokens to generate in the completion.
            Defaults to 256.
        top_p (float, optional): Total probability mass of tokens to consider at each step.
            Defaults to 1.
        top_k (int, optional): The cutoff where the model no longer selects the words.
            Defaults to 256.
        stop_sequences (List[str], optional): If the model encounters a stop sequence, it stops
            generating further tokens. Defaults to an empty list.
        session (Any, optional): A bedrock session. If provided, a new bedrock client will be
            created using this session. Defaults to None.
        client (Any, optional): The bedrock session client. If unset, a new one is created with
            boto3. Defaults to None.
        max_content_size (Optional[int], optional): If using a fine-tuned model, set this to the
            maximum content size. Defaults to None.
        extra_parameters (Dict[str, Any], optional): Any extra parameters to add to the request
            body (e.g., countPenalty for a21 models). Defaults to an empty dictionary.
        initial_rate_limit (int, optional): The initial internal rate limit in allowed requests
            per second for making LLM calls. This limit adjusts dynamically based on rate
            limit errors. Defaults to 5.
        timeout (int, optional): The timeout for completion requests in seconds. Defaults to 120.

    Example:
        .. code-block:: python

            # configure your AWS credentials using the AWS CLI
            # https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html

            from phoenix.evals import BedrockModel
            model = BedrockModel(model_id="anthropic.claude-v2")
    """

    model_id: str = "anthropic.claude-v2"
    temperature: float = 0.0
    max_tokens: int = 1024
    top_p: float = 1
    top_k: int = 256
    stop_sequences: List[str] = field(default_factory=list)
    session: Any = None
    client: Any = None
    max_content_size: Optional[int] = None
    extra_parameters: Dict[str, Any] = field(default_factory=dict)
    initial_rate_limit: int = 5
    timeout: int = 120

    def __post_init__(self) -> None:
        self._init_client()
        self._init_rate_limiter()

    @property
    def _model_name(self) -> str:
        return self.model_id

    def _init_client(self) -> None:
        if not self.client:
            if self.session:
                self.client = self.session.client("bedrock-runtime")
            else:
                try:
                    import boto3  # type:ignore

                    self.client = boto3.client("bedrock-runtime")
                except ImportError:
                    self._raise_import_error(
                        package_name="boto3",
                        package_min_version=MINIMUM_BOTO_VERSION,
                    )

    def _init_rate_limiter(self) -> None:
        self._rate_limiter = RateLimiter(
            rate_limit_error=self.client.exceptions.ThrottlingException,
            max_rate_limit_retries=10,
            initial_per_second_request_rate=self.initial_rate_limit,
            enforcement_window_minutes=1,
        )

    def _generate(self, prompt: Union[str, MultimodalPrompt], **kwargs: Dict[str, Any]) -> str:
        if isinstance(prompt, str):
            prompt = MultimodalPrompt.from_string(prompt)

        body = json.dumps(self._create_request_body(prompt))
        accept = "application/json"
        contentType = "application/json"

        response = self._rate_limited_completion(
            body=body, modelId=self.model_id, accept=accept, contentType=contentType
        )

        return self._parse_output(response) or ""

    async def _async_generate(
        self, prompt: Union[str, MultimodalPrompt], **kwargs: Dict[str, Any]
    ) -> str:
        if isinstance(prompt, str):
            prompt = MultimodalPrompt.from_string(prompt)

        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, partial(self._generate, prompt, **kwargs))

    def _rate_limited_completion(self, **kwargs: Any) -> Any:
        """Use tenacity to retry the completion call."""

        @self._rate_limiter.limit
        def _completion(**kwargs: Any) -> Any:
            try:
                return self.client.invoke_model(**kwargs)
            except Exception as e:
                exception_message = e.args[0]
                if not exception_message:
                    raise e

                if "Input is too long" in exception_message:
                    # Error from Anthropic models
                    raise PhoenixContextLimitExceeded(exception_message) from e
                elif "expected maxLength" in exception_message:
                    # Error from Titan models
                    raise PhoenixContextLimitExceeded(exception_message) from e
                elif "Prompt has too many tokens" in exception_message:
                    # Error from AI21 models
                    raise PhoenixContextLimitExceeded(exception_message) from e
                raise e

        return _completion(**kwargs)

    def _format_prompt_for_claude(self, prompt: str) -> List[Dict[str, str]]:
        # Claude requires prompt in the format of Human: ... Assisatnt:
        return [
            {"role": "user", "content": prompt},
        ]

    def _create_request_body(self, prompt: MultimodalPrompt) -> Dict[str, Any]:
        # The request formats for bedrock models differ
        # see https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html

        # TODO: Migrate to using the bedrock `converse` API
        # https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-call.html

        prompt_str = prompt.to_text_only_prompt()

        if self.model_id.startswith("ai21"):
            return {
                **{
                    "prompt": prompt_str,
                    "temperature": self.temperature,
                    "topP": self.top_p,
                    "maxTokens": self.max_tokens,
                    "stopSequences": self.stop_sequences,
                },
                **self.extra_parameters,
            }
        elif self.model_id.startswith("anthropic"):
            return {
                **{
                    "anthropic_version": "bedrock-2023-05-31",
                    "messages": self._format_prompt_for_claude(prompt_str),
                    "temperature": self.temperature,
                    "top_p": self.top_p,
                    "top_k": self.top_k,
                    "max_tokens": self.max_tokens,
                    "stop_sequences": self.stop_sequences,
                },
                **self.extra_parameters,
            }
        elif self.model_id.startswith("mistral"):
            return {
                **{
                    "prompt": prompt_str,
                    "max_tokens": self.max_tokens,
                    "temperature": self.temperature,
                    "stop": self.stop_sequences,
                    "top_p": self.top_p,
                    "top_k": self.top_k,
                },
                **self.extra_parameters,
            }
        else:
            if not self.model_id.startswith("amazon"):
                logger.warn(f"Unknown format for model {self.model_id}, returning titan format...")
            return {
                **{
                    "inputText": prompt_str,
                    "textGenerationConfig": {
                        "temperature": self.temperature,
                        "topP": self.top_p,
                        "maxTokenCount": self.max_tokens,
                        "stopSequences": self.stop_sequences,
                    },
                },
                **self.extra_parameters,
            }

    def _parse_output(self, response: Any) -> Any:
        if self.model_id.startswith("ai21"):
            body = json.loads(response.get("body").read())
            return body.get("completions")[0].get("data").get("text")
        elif self.model_id.startswith("anthropic"):
            body = json.loads(response.get("body").read().decode())
            return body.get("content")[0]["text"]
        elif self.model_id.startswith("amazon"):
            body = json.loads(response.get("body").read())
            return body.get("results")[0].get("outputText")
        elif self.model_id.startswith("mistral"):
            body = json.loads(response.get("body").read())
            return body.get("outputs")[0].get("text")
        else:
            body = json.loads(response.get("body").read())
            return body.get("results")[0].get("data").get("outputText")
