import logging
from typing import Any, List, Optional, Union

from langchain_core._api.beta_decorator import beta
from langchain_core.messages.ai import AIMessage
from langchain_core.prompt_values import ChatPromptValue
from langchain_core.runnables import Runnable
from langchain_core.runnables.config import RunnableConfig
from pydantic import ConfigDict

logger = logging.getLogger(__name__)


@beta(message="This API is in beta and can change in future.")
class AmazonQ(
    Runnable[Union[str, ChatPromptValue, List[ChatPromptValue]], ChatPromptValue]
):
    """Amazon Q Runnable wrapper.

    To authenticate, the AWS client uses the following methods to
    automatically load credentials:
    https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html

    Make sure the credentials / roles used have the required policies to
    access the Amazon Q service.

    """

    region_name: Optional[str] = None
    """AWS region name. If not provided, will be extracted from environment."""

    credentials: Optional[Any] = None
    """Amazon Q credentials used to instantiate the client if the client is not provided."""  # noqa: E501

    client: Optional[Any] = None
    """Amazon Q client."""

    application_id: Optional[str] = None

    """Store the full response from Amazon Q."""

    parent_message_id: Optional[str] = None

    conversation_id: Optional[str] = None

    chat_mode: str = "RETRIEVAL_MODE"

    model_config = ConfigDict(
        extra="forbid",
    )

    def __init__(
        self,
        region_name: Optional[str] = None,
        credentials: Optional[Any] = None,
        client: Optional[Any] = None,
        application_id: Optional[str] = None,
        parent_message_id: Optional[str] = None,
        conversation_id: Optional[str] = None,
        chat_mode: str = "RETRIEVAL_MODE",
    ):
        self.region_name = region_name
        self.credentials = credentials
        self.client = client or self.validate_environment()
        self.application_id = application_id
        self.parent_message_id = parent_message_id
        self.conversation_id = conversation_id
        self.chat_mode = chat_mode

    def invoke(
        self,
        input: Union[str, ChatPromptValue, List[ChatPromptValue]],
        config: Optional[RunnableConfig] = None,
        **kwargs: Any,
    ) -> ChatPromptValue:
        """Call out to Amazon Q service.

        Args:
            input: The prompt to pass into the model.

        Returns:
            The string generated by the model.

        Example:
            ```python
            model = AmazonQ(
                credentials=your_credentials,
                application_id=your_app_id
            )
            response = model.invoke("Tell me a joke")
            ```
        """
        try:
            # Prepare the request
            request = {
                "applicationId": self.application_id,
                "userMessage": self.convert_langchain_messages_to_q_input(
                    input
                ),  # Langchain's input comes in the form of an array of "messages".
                # We must convert to a single string for Amazon Q's use
                "chatMode": self.chat_mode,
            }
            if self.conversation_id:
                request.update(
                    {
                        "conversationId": self.conversation_id,
                        "parentMessageId": self.parent_message_id,
                    }
                )

            # Call Amazon Q
            if self.client is None:
                raise ValueError("Amazon Q client is not initialized")
            response = self.client.chat_sync(**request)

            # Extract the response text
            if "systemMessage" in response:
                ai_message = AIMessage(
                    content=response["systemMessage"], response_metadata=response
                )
                return ChatPromptValue(messages=[ai_message])
            else:
                raise ValueError("Unexpected response format from Amazon Q")

        except Exception as e:
            if "Prompt Length" in str(e):
                # Convert input to string to get its length
                input_str = self.convert_langchain_messages_to_q_input(input)
                logger.info(f"Prompt Length: {len(input_str)}")
                logger.info(f"""Prompt:
                {input_str}""")
            raise ValueError(f"Error raised by Amazon Q service: {e}")

    def validate_environment(self) -> Any:
        """Don't do anything if client provided externally"""
        # If the client is not provided, and the user_id is not provided in the class
        # constructor, throw an error saying one or the other needs to be provided
        if self.credentials is None:
            raise ValueError(
                "Either the credentials or the client needs to be provided."
            )

        """Validate that AWS credentials to and python package exists in environment."""
        try:
            import boto3

            try:
                if self.region_name is not None:
                    client = boto3.client(
                        "qbusiness", self.region_name, **self.credentials
                    )
                else:
                    # use default region
                    client = boto3.client("qbusiness", **self.credentials)

            except Exception as e:
                raise ValueError(
                    "Could not load credentials to authenticate with AWS client. "
                    "Please check that credentials in the specified "
                    "profile name are valid."
                ) from e

        except ImportError:
            raise ImportError(
                "Could not import boto3 python package. "
                "Please install it with `pip install boto3`."
            )
        return client

    def convert_langchain_messages_to_q_input(
        self, input: Union[str, ChatPromptValue, List[ChatPromptValue]]
    ) -> str:
        # If it is just a string and not a ChatPromptTemplate collection just
        # return string
        if isinstance(input, str):
            return input
        elif isinstance(input, ChatPromptValue):
            return input.to_string()
        elif isinstance(input, list):
            # For list of ChatPromptValue, concatenate their string representations
            return "\n".join(item.to_string() for item in input)
        else:
            # This should never happen given the type annotation, but just in case
            return str(input)
