from typing import Dict, List, Optional, Union

from synthetic_agents.common.functions import safe_placeholder_mapping
from synthetic_agents.prompt.loader import (
    PromptTemplateLoader,
    TextPromptTemplateLoader,
)


class PromptBuilder:
    """
    This class is used to build prompts to be passed to a language model by replacing placeholders
    in a prompt template with actual values.
    """

    def __init__(
        self,
        prompt_template_loader: PromptTemplateLoader,
        placeholder_values: Optional[Dict[str, Union[str, int, float]]] = None,
    ):
        """
        Creates a prompt builder.

        :param prompt_template_loader: an object that reads the prompt template from some source.
        :param placeholder_values: a dictionary of placeholder keys and values to replace the
            placeholders in the template.
        """
        self.prompt_template_loader = prompt_template_loader
        self.placeholder_values = placeholder_values

    def build_prompt(self) -> str:
        """
        Builds a prompt by populating placeholders in a template.

        :return: prompt template with some or all of the placeholders filled.
        """
        if not self.prompt_template_loader.template:
            self.prompt_template_loader.initialize_loader()

        if not self.placeholder_values:
            return self.prompt_template_loader.template

        # The prompt template is a text with placeholders surrounded by curly brackets. For
        # instance, {demographics}. So, we could fill these placeholders with the string.format
        # function and an unpacked dictionary if all keys in the template were present in the
        # placeholder_values. Since that may not always be the case, we use a helper function that
        # preserves placeholders not part of the placeholder_values dictionary.
        return safe_placeholder_mapping(
            template=self.prompt_template_loader.template,
            placeholder_values=self.placeholder_values,
        )


class FewShotPromptBuilder(PromptBuilder):
    """
    This class is used to build prompts with few-shot examples to be passed to a language model by
    replacing placeholders in a prompt template with actual values. The template is formed by
    four main sections:
    1. Instruction
    2. Examples (each example has an input and an output)
    3. Input
    4. Output
    """

    _PROMPT_TEMPLATE = """
INSTRUCTIONS:
{instructions}

{examples}

INPUT:
{true_input}

OUTPUT:
    """

    _EXAMPLE_PROMPT_TEMPLATE = """
INPUT EXAMPLE {example_number}:
{input_example}

OUTPUT EXAMPLE {example_number}:
{output_example}
        """

    def __init__(self, instructions: str, true_input: str):
        """
        Creates a few-shot prompt builder.

        :param instructions: a text with instructions to the be added to the prompt.
        :param true_input: a text with the true input to be passed to the prompt. The one we want
            to generate text to.
        """
        super().__init__(
            prompt_template_loader=TextPromptTemplateLoader(FewShotPromptBuilder._PROMPT_TEMPLATE),
            placeholder_values={
                "instructions": instructions,
                "examples": "",
                "true_input": true_input,
            },
        )

        self._few_shot_example_prompts: List[str] = []

    def add_example(self, example_input: str, example_output: str):
        """
        Adds a prompt to the list of examples. Each example has an input and an expected output.

        :param example_input: input to be used as an example.
        :param example_output: expected output for the example input.
        """

        self._few_shot_example_prompts.append(
            safe_placeholder_mapping(
                FewShotPromptBuilder._EXAMPLE_PROMPT_TEMPLATE,
                {
                    "example_number": len(self._few_shot_example_prompts) + 1,
                    "input_example": example_input,
                    "output_example": example_output,
                },
            )
        )

        # Update the {examples} placeholder with the full list of examples.
        self.placeholder_values["examples"] = "\n".join(self._few_shot_example_prompts)
