from dapr_agents.prompt.utils.jinja import (
    render_jinja_template,
    extract_jinja_variables,
)
from dapr_agents.prompt.utils.fstring import (
    render_fstring_template,
    extract_fstring_variables,
)
from dapr_agents.prompt.utils.chat import ChatPromptHelper
from dapr_agents.prompt.base import PromptTemplateBase
from dapr_agents.types.message import (
    BaseMessage,
    SystemMessage,
    UserMessage,
    AssistantMessage,
    ToolMessage,
    MessagePlaceHolder,
)
from typing import Any, Dict, List, Tuple, Union, Literal, Optional
from pydantic import Field
import logging

logger = logging.getLogger(__name__)

DEFAULT_FORMATTER_MAPPING = {
    "f-string": render_fstring_template,
    "jinja2": render_jinja_template,
}

DEFAULT_VARIABLE_EXTRACTOR_MAPPING = {
    "f-string": extract_fstring_variables,
    "jinja2": extract_jinja_variables,
}


class ChatPromptTemplate(PromptTemplateBase):
    """
    A template class designed to handle chat-based prompts. This class can format a sequence of chat messages
    and merge chat history with provided variables or placeholders.

    Attributes:
        messages (List[Union[Tuple[str, str], Dict[str, Any], BaseMessage, MessagePlaceHolder]]): A list of messages that make up the prompt template.
        template_format (Literal["f-string", "jinja2"]): The format used for rendering the template.
    """

    messages: List[
        Union[Tuple[str, str], Dict[str, Any], BaseMessage, MessagePlaceHolder]
    ] = Field(default_factory=list)
    template_format: Literal["f-string", "jinja2"] = "f-string"

    _ROLE_MAP = {
        "system": SystemMessage,
        "user": UserMessage,
        "assistant": AssistantMessage,
        "tool": ToolMessage,
    }

    def format_prompt(
        self, template_format: Optional[str] = None, **kwargs: Any
    ) -> List[Dict[str, Any]]:
        """
        Format the prompt by processing placeholders, rendering the template with variables,
        and then incorporating both plain text and role-based messages, if applicable.

        Returns:
            List[Dict[str, Any]]: The list of formatted messages as dictionaries.
        """
        template_format = template_format or self.template_format
        all_variables = self.prepare_variables_for_formatting(**kwargs)

        # Check for undeclared or missing variables
        extra_variables = [
            var
            for var in all_variables
            if var not in self.input_variables and var not in self.pre_filled_variables
        ]
        if extra_variables:
            raise ValueError(f"Undeclared variables were passed: {extra_variables}")

        missing_variables = [
            var for var in self.input_variables if var not in all_variables
        ]
        if missing_variables:
            logger.info(f"Some input variables were not provided: {missing_variables}")

        rendered_messages = []

        for item in self.messages:
            # Process MessagePlaceHolder with dynamic messages from all_variables
            if isinstance(item, MessagePlaceHolder):
                variable_name = item.variable_name
                if variable_name in all_variables:
                    normalized_messages = ChatPromptHelper.normalize_chat_messages(
                        all_variables[variable_name]
                    )
                    rendered_messages.extend(
                        [msg.model_dump() for msg in normalized_messages]
                    )
                else:
                    logger.info(
                        f"MessagePlaceHolder variable '{variable_name}' was not provided."
                    )

            # Process BaseMessage, Tuple, and Dict with parse_as_messages
            else:
                role, content = ChatPromptHelper.extract_role_and_content(item)
                formatted_content = ChatPromptHelper.format_content(
                    content, template_format, **all_variables
                )
                parsed_messages, plain_text = ChatPromptHelper.parse_as_messages(
                    formatted_content
                )

                # Add the plain text only if parsed messages are also returned
                if plain_text and parsed_messages:
                    rendered_messages.append(
                        ChatPromptHelper.create_message(
                            role, plain_text, {}
                        ).model_dump()
                    )

                # Add parsed role-based messages if they exist
                if parsed_messages:
                    rendered_messages.extend(
                        [msg.model_dump() for msg in parsed_messages]
                    )
                else:
                    # If only plain text is present (no parsed messages), add it as a single message
                    rendered_messages.append(
                        ChatPromptHelper.create_message(
                            role, plain_text or formatted_content, {}
                        ).model_dump()
                    )

        return rendered_messages

    @classmethod
    def from_messages(
        cls,
        messages: List[
            Union[Tuple[str, str], Dict[str, Any], BaseMessage, MessagePlaceHolder]
        ],
        template_format: str = "f-string",
    ) -> "ChatPromptTemplate":
        """
        Create a ChatPromptTemplate from a list of messages, including placeholders.

        Args:
            messages (List[Union[Tuple[str, str], Dict[str, Any], BaseMessage, MessagePlaceHolder]]):
                The list of messages that define the template.
            template_format (str): The format of the template, either "f-string" or "jinja2". Default is "f-string".

        Returns:
            ChatPromptTemplate: A new instance of the template with extracted input variables.
        """
        input_vars = set()

        for msg in messages:
            content = None  # Initialize content to None

            # Handle MessagePlaceHolder by adding its variable directly to input_vars
            if isinstance(msg, MessagePlaceHolder):
                input_vars.add(msg.variable_name)

            elif isinstance(msg, tuple) and len(msg) == 2:
                content = msg[1]

            elif isinstance(msg, dict) and "content" in msg:
                content = msg["content"]

            elif isinstance(msg, BaseMessage):
                content = msg.content

            if isinstance(content, str):
                input_vars.update(
                    DEFAULT_VARIABLE_EXTRACTOR_MAPPING[template_format](content)
                )

        return cls(
            input_variables=list(input_vars),
            messages=messages,
            template_format=template_format,
        )