from floki.prompt.utils.jinja import render_jinja_template, extract_jinja_variables
from floki.prompt.utils.fstring import render_fstring_template, extract_fstring_variables
from floki.prompt.utils.chat import ChatPromptHelper
from floki.prompt.base import PromptTemplateBase
from floki.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)