import streamlit as st
from sqlalchemy import select

from synthetic_agents.database.config import get_db
from synthetic_agents.database.entity.prompt_template import PromptTemplate
from synthetic_agents.webapp.common.constants import APPLICATION_TYPES
from synthetic_agents.webapp.common.functions import (
    read_available_agent_types,
    read_available_prompt_templates,
    wait_and_refresh,
)

DEFAULT_PLACEHOLDER_INFO = """
:blue[The following placeholders will be automatically populated by the language model if present
 in the prompt:]

:blue[**{all_memories}**: life and chat memories retrieved by the working memory.]\n
:blue[**{life_memories}**: life memories retrieved by the working memory.]\n
:blue[**{chat_memories}**: chat messages retrieved by the working memory.]\n
:blue[**{current_date}**: the current date and time of the chat.]\n
:blue[**{chat_history}**: the most recent chat messages.]\n

:blue[The number of chat messages retrieved by the working memory and the maximum length of the
 chat history are determined at the chat creation. In addition, dictionary keys defining agent's
 attributes can also be used as placeholders which will be replaced with the corresponding values
 during the chat.]
"""


class PromptTemplateBuilderComponent:
    """
    This class handles the creation of prompt templates for user and coach agents.
    """

    def __init__(self, component_key: str):
        """
        Creates the prompt template builder component.

        :param component_key: unique identifier for the component in a page.
        """
        self.component_key = component_key

    def create_component(self):
        """
        Creates all the widgets and logic of the prompt template builder component.
        """
        st.header("Prompt Template")
        col1, col2 = st.columns(2)
        with col1:
            application_type = st.selectbox(
                "Application Type*",
                key=f"{self.component_key}_application_type_selector",
                options=sorted(APPLICATION_TYPES),
            )
        with col2:
            agent_type = st.selectbox(
                "Agent Type*",
                key=f"{self.component_key}_agent_type_selector",
                options=read_available_agent_types(),
                format_func=lambda type_: type_.description,
            )
        prompt_template = st.selectbox(
            "Version*",
            key=f"{self.component_key}_prompt_template_version_selector",
            options=read_available_prompt_templates(
                agent_type=agent_type.type_id if agent_type else None,
                application_type=application_type,
            ),
            format_func=lambda template: f"{template.version} - "
            f"{template.content[:min(len(template.content), 30)]}...",
        )

        st.markdown(DEFAULT_PLACEHOLDER_INFO)

        template_content = st.text_area(
            "Content",
            key=f"{self.component_key}_template_content_text_area",
            value=prompt_template.content if prompt_template else "",
            height=400,
        )

        new_version = st.text_input(
            "New Version",
            key=f"{self.component_key}_new_version_input",
            placeholder="Enter a version for the new template (e.g., 0.0.2).",
        )

        col1, col2, col3, _ = st.columns([0.1, 0.1, 0.1, 0.7])
        with col1:
            new_template_button = st.button("New", key=f"{self.component_key}_new_button")
        with col2:
            edit_template_button = st.button(
                "Edit", key=f"{self.component_key}_edit_button", disabled=prompt_template is None
            )
        with col3:
            delete_template_button = st.button(
                "Delete",
                key=f"{self.component_key}_delete_button",
                disabled=prompt_template is None,
            )

        if new_template_button:
            if PromptTemplateBuilderComponent._exists_in_database(
                application_type=application_type,
                agent_type=agent_type.type_id,
                template_version=new_version,
            ):
                st.error(
                    "There is already a template with the same version for the application "
                    "and agent types selected."
                )
            else:
                # There's no other template with the same version for the within the
                # application type and agent type.
                new_template = PromptTemplate(
                    application_type=application_type,
                    agent_type=agent_type.type_id,
                    version=new_version,
                    content=template_content,
                )
                PromptTemplateBuilderComponent._update_or_delete_prompt_template(
                    new_template, delete=False
                )
                st.success("Template created successfully!")
                read_available_prompt_templates.clear()
                wait_and_refresh()

        elif edit_template_button:
            prompt_template.content = template_content
            PromptTemplateBuilderComponent._update_or_delete_prompt_template(
                prompt_template, delete=False
            )
            st.success("Template edited successfully!")
            read_available_prompt_templates.clear()
            wait_and_refresh()
        elif delete_template_button:
            PromptTemplateBuilderComponent._update_or_delete_prompt_template(
                prompt_template, delete=True
            )
            st.success("Template deleted successfully!")
            read_available_prompt_templates.clear()
            wait_and_refresh()

    @staticmethod
    def _exists_in_database(application_type: str, agent_type: str, template_version: str):
        """
        Checks whether a prompt template with the same version exists in the database for an
        application and agent types.

        :param application_type: type of the application.
        :param agent_type: type of the agent.
        :param template_version: version of the prompt template.
        :return: True if it exists in the database, False otherwise.
        """
        db = next(get_db())
        existing_template = db.execute(
            select(PromptTemplate).where(
                PromptTemplate.application_type == application_type,
                PromptTemplate.agent_type == agent_type,
                PromptTemplate.version == template_version,
            )
        ).first()

        return existing_template is not None

    @staticmethod
    def _update_or_delete_prompt_template(prompt_template: PromptTemplate, delete: bool):
        """
        Adds, edits or deletes a prompt template in the database.

        :param prompt_template: a prompt template persistent object.
        :param delete: whether we want to delete an existing prompt template.
        """
        db = next(get_db())
        if delete:
            db.delete(prompt_template)
        else:
            db.add(prompt_template)
        db.commit()
