# -*- coding: utf-8 -*-
"""
Templates for text generated by the bot.  There are two types:
 - templates used to generate prompts for the AI
 - templates used to generate UI messages for the user
"""

import enum
import functools
import textwrap
import typing


@functools.total_ordering
class Templates(enum.Enum):
    """
    Enumeration of all different templates.
    """

    COMMAND_LOBOTOMIZE_RESPONSE = "command_lobotomize_response"

    IMAGE_DETACH = "image_detach"
    IMAGE_CONFIRMATION = "image_confirmation"
    IMAGE_GENERATION_ERROR = "image_generation_error"
    IMAGE_UNAUTHORIZED = "image_unauthorized"

    # prompts to the AI to generate text responses
    PROMPT = "prompt"
    PROMPT_HISTORY_LINE = "prompt_history_line"
    PROMPT_IMAGE_COMING = "prompt_image_coming"

    # prompts to the AI to generate image keywords
    PROMPT_IMAGE_KEYWORDS = "prompt_image_keywords"

    def __str__(self) -> str:
        return self.value

    def __lt__(self, other: "Templates") -> bool:
        return str(self.value) < str(other.value)


class TemplateToken(str, enum.Enum):
    """
    Enumeration of all tokens used in templates.
    Tokens are variable substitutions into the templates.
    """

    AI_NAME = "AI_NAME"
    PERSONA = "PERSONA"
    IMAGE_COMING = "IMAGE_COMING"
    IMAGE_PROMPT = "IMAGE_PROMPT"
    MESSAGE_HISTORY = "MESSAGE_HISTORY"
    USER_MESSAGE = "USER_MESSAGE"
    USER_NAME = "USER_NAME"

    def __str__(self):
        return "{" + self.value + "}"


class TemplateStore:
    """
    Data object storing all template definitions and default values.
    """

    # mapping of template names to tokens allowed in that template
    #  key: template name
    #  value: tuple of (list of tokens, description, is_an_ai_prompt)
    TEMPLATES: typing.Dict[
        Templates, typing.Tuple[typing.List[TemplateToken], str, bool]
    ] = {
        Templates.COMMAND_LOBOTOMIZE_RESPONSE: (
            [
                TemplateToken.AI_NAME,
                TemplateToken.USER_NAME,
            ],
            "Displayed in Discord after a successful /lobotomize command.  "
            + "Both the discord users and the bot AI will see this message.",
            False,
        ),
        Templates.PROMPT: (
            [
                TemplateToken.AI_NAME,
                TemplateToken.IMAGE_COMING,
                TemplateToken.MESSAGE_HISTORY,
                TemplateToken.PERSONA,
            ],
            "The main prompt sent to Oobabooga to generate a response from "
            + "the bot AI.  The AI's reply to this prompt will be sent to "
            + "discord as the bot's response.",
            True,
        ),
        Templates.PROMPT_HISTORY_LINE: (
            [
                TemplateToken.USER_MESSAGE,
                TemplateToken.USER_NAME,
            ],
            "Part of the AI response-generation prompt, this is used to "
            + "render a single line of chat history.  A list of these, "
            + "one for each past chat message, will become {MESSAGE_HISTORY} "
            + "and inserted into the main prompt",
            True,
        ),
        Templates.PROMPT_IMAGE_COMING: (
            [
                TemplateToken.AI_NAME,
            ],
            "Part of the AI response-generation prompt, this is used to "
            + "inform the AI that it is in the process of generating an "
            + "image.",
            True,
        ),
        Templates.PROMPT_IMAGE_KEYWORDS: (
            [
                TemplateToken.AI_NAME,
                TemplateToken.USER_MESSAGE,
            ],
            """
            Sent to Oobabooga, along with the user's image request, to generate
            image keywords.  The AI's response to this prompt will then be sent
            to Stable Diffusion to generate an image.
            """,
            True,
        ),
        Templates.IMAGE_DETACH: (
            [
                TemplateToken.IMAGE_PROMPT,
                TemplateToken.USER_NAME,
            ],
            "Shown in Discord when the user selects to discard an image "
            + "that Stable Diffusion had generated.",
            False,
        ),
        Templates.IMAGE_CONFIRMATION: (
            [
                TemplateToken.IMAGE_PROMPT,
                TemplateToken.USER_NAME,
            ],
            "Shown in Discord when an image is first generated from "
            + "Stable Diffusion.  This should prompt the user to either "
            + "save or discard the image.",
            False,
        ),
        Templates.IMAGE_GENERATION_ERROR: (
            [
                TemplateToken.IMAGE_PROMPT,
                TemplateToken.USER_NAME,
            ],
            "Shown in Discord when the we could not contact Stable Diffusion "
            + "to generate an image.",
            False,
        ),
        Templates.IMAGE_UNAUTHORIZED: (
            [TemplateToken.USER_NAME],
            "Shown in Discord privately to a user if they try to regenerate "
            "an image that was requested by someone else.",
            False,
        ),
    }

    DEFAULT_TEMPLATES: typing.Dict[Templates, str] = {
        Templates.PROMPT: textwrap.dedent(
            """
            You are in a chat room with multiple participants.
            Below is a transcript of recent messages in the conversation.
            Write the next one to three messages that you would send in this
            conversation, from the point of view of the participant named
            {AI_NAME}.

            {PERSONA}

            All responses you write must be from the point of view of
            {AI_NAME}.
            ### Transcript:
            {MESSAGE_HISTORY}
            {IMAGE_COMING}
            """
        ),
        Templates.PROMPT_HISTORY_LINE: textwrap.dedent(
            """
            {USER_NAME} says:
            {USER_MESSAGE}

            """
        ),
        Templates.PROMPT_IMAGE_COMING: textwrap.dedent(
            """
            {AI_NAME}: is currently generating an image, as requested.
            """
        ),
        Templates.PROMPT_IMAGE_KEYWORDS: textwrap.dedent(
            """
            Below is an instruction that describes a task, paired with an input
            that provides further context. Write a response that appropriately
            completes the request.

            ### Instruction:
            The input text asks for an image to be created.  Please create a list
            of keywords to describe that image.

            ### Input:
            {USER_MESSAGE}

            ### Response:
            """
        ),
        Templates.IMAGE_DETACH: textwrap.dedent(
            """
            {USER_NAME} asked for an image with the prompt:
                '{IMAGE_PROMPT}'
            ...but couldn't find a suitable one.
            """
        ),
        Templates.IMAGE_CONFIRMATION: textwrap.dedent(
            """
            {USER_NAME}, is this what you wanted?
            If no choice is made, this message will 💣 self-destruct 💣 in 3 minutes.
            """
        ),
        Templates.IMAGE_GENERATION_ERROR: textwrap.dedent(
            """
            Something went wrong generating your image.  Sorry about that!
            """
        ),
        Templates.IMAGE_UNAUTHORIZED: textwrap.dedent(
            """
            Sorry, only {USER_NAME} can press the buttons.
            """
        ),
        Templates.COMMAND_LOBOTOMIZE_RESPONSE: textwrap.dedent(
            """
            Ummmm... what were we talking about?
            """
        ),
    }

    def __init__(self, settings: dict):
        self.templates: typing.Dict[Templates, TemplateMessageFormatter] = {}
        for template, (tokens, purpose, is_ai_prompt) in self.TEMPLATES.items():
            template_name = str(template)
            template_fmt = settings[template_name]
            if template_fmt is None:
                raise ValueError(f"Template {template_name} has no default format")
            self.add_template(template, template_fmt, tokens, purpose, is_ai_prompt)

    def add_template(
        self,
        template_name: Templates,
        format_str: str,
        allowed_tokens: typing.List[TemplateToken],
        purpose: str,
        is_ai_prompt: bool,
    ):
        self.templates[template_name] = TemplateMessageFormatter(
            template_name,
            format_str,
            allowed_tokens,
            purpose,
            is_ai_prompt,
        )

    def format(
        self, template_name: Templates, format_args: typing.Dict[TemplateToken, str]
    ) -> str:
        return self.templates[template_name].format(format_args)


class TemplateMessageFormatter:
    """
    Validates that templates are safe to run string.format() on, and
    runs string.format()
    """

    def __init__(
        self,
        template_name: Templates,
        template: str,
        allowed_tokens: typing.List[TemplateToken],
        purpose: str,
        is_ai_prompt: bool,
    ):
        self._validate_format_string(template_name, template, allowed_tokens)
        self.template_name = template_name
        self.template = template
        self.allowed_tokens = allowed_tokens
        self.purpose = purpose
        self.is_ai_prompt = is_ai_prompt

    def __str__(self):
        return self.template

    def format(self, format_args: typing.Dict[TemplateToken, str]) -> str:
        return self.template.format(**format_args)

    @staticmethod
    def _validate_format_string(
        template_name: Templates,
        format_str: str,
        allowed_args: typing.List[TemplateToken],
    ):
        def find_all_ch(string: str, char: str) -> typing.Generator[int, None, None]:
            # find all indices of ch in s
            for i, letter in enumerate(string):
                if letter == char:
                    yield i

        # raises if fmt_string contains any args not in allowed_args
        allowed_close_brace_indices: typing.Set[int] = set()

        for open_brace_idx in find_all_ch(format_str, "{"):
            for allowed_arg in allowed_args:
                idx_end = open_brace_idx + len(allowed_arg) + 1
                next_substr = format_str[open_brace_idx : idx_end + 1]
                if next_substr == "{" + allowed_arg + "}":
                    allowed_close_brace_indices.add(idx_end)
                    break
            else:
                raise ValueError(
                    f"invalid template: {template_name} contains "
                    + f"an argument not in {allowed_args}"
                )
        for close_brace_idx in find_all_ch(format_str, "}"):
            if close_brace_idx not in allowed_close_brace_indices:
                raise ValueError(
                    f"invalid template: {template_name} contains "
                    + f"an argument not in {allowed_args}"
                )
