from dapr_agents.prompt.base import PromptTemplateBase
from dapr_agents.prompt.utils.string import StringPromptHelper
from typing import Any, Union, Literal


class StringPromptTemplate(PromptTemplateBase):
    """
    A template class designed to handle string-based prompts. This class can format a string template
    by replacing variables and merging additional context or keywords.

    Attributes:
        template (str): The template string that defines the prompt structure.
        template_format (Literal["f-string", "jinja2"]): The format used for rendering the template, either f-string or Jinja2.
    """

    template: str
    template_format: Literal["f-string", "jinja2"] = "f-string"

    def format_prompt(self, **kwargs: Any) -> str:
        """
        Format the prompt by replacing variables in the template, using any provided keyword arguments.
        Validates the template structure before formatting.

        Args:
            **kwargs: Additional keyword arguments to replace variables in the template.

        Returns:
            str: The formatted prompt with all variables replaced.

        Raises:
            ValueError: If required variables are missing or extra undeclared variables are passed.
        """
        # Extract the required input variables from the template
        input_variables = StringPromptHelper.extract_variables(
            self.template, self.template_format
        )

        # Check for missing variables
        missing_vars = [var for var in input_variables if var not in kwargs]
        if missing_vars:
            raise ValueError(f"Missing required variables in template: {missing_vars}")

        # Check for extra variables that were not expected
        extra_variables = [var for var in kwargs if var not in input_variables]
        if extra_variables:
            raise ValueError(f"Undeclared variables were passed: {extra_variables}")

        # Prepare the variables for formatting
        kwargs = self.prepare_variables_for_formatting(**kwargs)

        # Use the helper to format the content
        return StringPromptHelper.format_content(
            self.template, self.template_format, **kwargs
        )

    @classmethod
    def from_template(
        cls, template: str, template_format: str = "f-string", **kwargs: Any
    ) -> "StringPromptTemplate":
        """
        Create a StringPromptTemplate from a template string.

        Args:
            template (str): The template string that defines the structure of the prompt.
            template_format (str): The format of the template, either "f-string" or "jinja2". Default is "f-string".
            **kwargs: Additional keyword arguments for the constructor.

        Returns:
            StringPromptTemplate: A new instance of the template with extracted input variables.
        """
        input_variables = StringPromptHelper.extract_variables(
            template, template_format
        )
        return cls(
            template=template,
            template_format=template_format,
            input_variables=input_variables,
            **kwargs,
        )

    def __add__(
        self, other: Union[str, "StringPromptTemplate"]
    ) -> "StringPromptTemplate":
        """
        Override the + operator to allow for combining prompt templates.

        Args:
            other: Another prompt template or string to be combined with the current one.

        Returns:
            StringPromptTemplate: A new instance of the combined templates with merged variables.
        """
        if isinstance(other, StringPromptTemplate):
            # Ensure both templates use the same format
            if self.template_format != other.template_format:
                raise ValueError(
                    "Adding prompt templates only supported for the same template format."
                )

            # Combine input variables
            input_variables = list(
                set(self.input_variables) | set(other.input_variables)
            )

            # Combine template strings
            template = self.template + other.template

            # Combine pre-filled variables
            pre_filled_variables = {
                **self.pre_filled_variables,
                **other.pre_filled_variables,
            }

            # Create and return a new combined StringPromptTemplate
            return StringPromptTemplate(
                template=template,
                template_format=self.template_format,
                input_variables=input_variables,
                pre_filled_variables=pre_filled_variables,
            )
        elif isinstance(other, str):
            # If other is a string, convert it to a StringPromptTemplate and combine
            return self + StringPromptTemplate.from_template(
                other, template_format=self.template_format
            )
        else:
            # Raise error for unsupported types
            raise NotImplementedError(f"Unsupported operand type for +: {type(other)}")