import json

import streamlit as st

from synthetic_agents.generator.life_fact import LifeFactGenerator
from synthetic_agents.generator.user_profile import UserProfileGenerator
from synthetic_agents.model.agent import Agent
from synthetic_agents.model.memory import LifeMemory
from synthetic_agents.webapp.common.constants import (
    APPLICATION_TYPES,
    DEFAULT_RANDOM_AGENT_NUM_LIFE_FACTS,
    LLMS,
)
from synthetic_agents.webapp.common.functions import (
    read_available_agent_types,
    read_available_prompt_templates,
)
from synthetic_agents.webapp.session.agent_builder import AgentBuilderSession


class AgentBuilderComponent:
    """
    This class is responsible for handling agent creation for usage by a chat application.
    """

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

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

        # Values saved within page loading and available to the next widgets to be rendered.
        # Not persisted through the session.
        self._application_type = None
        self._agent_type = None
        self._prompt_template = None
        self._llm = None
        self._temperature = None

    def create_component(self):
        """
        Creates all the widgets and logic of the agent builder component.
        """
        if st.button("Refresh Prompt Template List", key=f"{self.component_key}_refresh_button"):
            read_available_prompt_templates.clear()

        self._create_header_widgets()

        self._llm = st.selectbox("Model", key=f"{self.component_key}_llm_selector", options=LLMS)
        self._temperature = st.slider(
            "Temperature",
            key=f"{self.component_key}_temperature_slider",
            min_value=0.0,
            max_value=1.0,
            value=0.75,
        )

        self._create_user_profile_area()
        self._create_life_facts_area()

        st.divider()

        if (
            self._application_type is not None
            and self._agent_type is not None
            and self._prompt_template is not None
        ):
            create_agent = st.button("Create Agent")
            st.write(
                "*:blue[Pressing the button above will persist the agent to the database and "
                "assign an ID to it. The new agent will then be available for usage in the "
                "chat application.]*"
            )
            if create_agent:
                new_agent_id = Agent.create(
                    application_type=self._application_type,
                    agent_type=self._agent_type.type_id,
                    prompt_template_version=self._prompt_template.version,
                    attributes=AgentBuilderSession.get_stored_user_profile(),
                    life_facts=AgentBuilderSession.get_stored_life_facts(),
                )
                st.success(f"Agent {new_agent_id} created successfully!")
        else:
            st.write(
                "*:blue[A button will show up here to create an agent when all the required "
                "information has been provided: application type, agent type, prompt "
                "template version, agent attributes and life facts (it can be empty).]*"
            )

    def _create_header_widgets(self):
        """
        Creates widgets for application, agent type and prompt template version selection.
        """
        col1, col2 = st.columns(2)
        with col1:
            self._application_type = st.selectbox(
                "Application Type*",
                key=f"{self.component_key}_application_type_selector",
                options=sorted(APPLICATION_TYPES),
            )
        with col2:
            self._agent_type = st.selectbox(
                "Agent Type*",
                key=f"{self.component_key}_agent_type_selector",
                options=read_available_agent_types(),
                format_func=lambda agent_type: agent_type.description,
            )
        self._prompt_template = st.selectbox(
            "Prompt Template Version*",
            key=f"{self.component_key}_prompt_template_version_selector",
            options=read_available_prompt_templates(
                agent_type=self._agent_type.type_id if self._agent_type else None,
                application_type=self._application_type,
            ),
            format_func=lambda template: f"{template.version} - "
            f"{template.content[:min(len(template.content), 30)]}...",
        )
        if self._prompt_template:
            with st.expander("Prompt Template", expanded=True):
                st.code(self._prompt_template.content, language=None)

    def _create_user_profile_area(self):
        user_profile = AgentBuilderSession.get_stored_user_profile()
        user_profile = st.text_area(
            "User Profile",
            key=f"{self.component_key}_user_profile_text_area",
            value=json.dumps(user_profile if user_profile is not None else {}, indent=4),
            height=400,
        )

        user_profile = json.loads(user_profile)
        AgentBuilderSession.store_user_profile(user_profile)
        if st.button(
            "Generate User Profile",
            key=f"{self.component_key}_generate_user_profile_button",
        ):
            with st.spinner():
                user_profile = UserProfileGenerator(
                    temperature=self._temperature, llm_name=self._llm
                ).generate(1)[0]
                AgentBuilderSession.store_user_profile(user_profile)
                st.rerun()

    def _create_life_facts_area(self):
        num_facts = st.number_input(
            "Number of Facts",
            key=f"{self.component_key}_number_life_facts_number_input",
            value=DEFAULT_RANDOM_AGENT_NUM_LIFE_FACTS,
        )
        life_facts = AgentBuilderSession.get_stored_life_facts()
        life_facts = life_facts if life_facts is not None else []
        life_facts = "\n".join(
            [
                json.dumps(m.to_simplified_json(["creation_timestamp", "content"]))
                for m in life_facts
            ]
        )
        life_facts = st.text_area(
            "Life Facts",
            key=f"{self.component_key}_life_facts_text_area",
            value=life_facts,
            height=400,
        )
        life_facts = [
            LifeMemory.from_json(json.loads(life_fact))
            for life_fact in life_facts.split("\n")
            if life_fact != ""
        ]
        AgentBuilderSession.store_generated_life_facts(life_facts)

        if st.button(
            "Generate Life Facts",
            key=f"{self.component_key}_generate_life_facts_button",
        ):
            with st.spinner():
                generator = LifeFactGenerator(
                    temperature=self._temperature,
                    user_profile=AgentBuilderSession.get_stored_user_profile(),
                    llm_name=self._llm,
                )
                life_facts = generator.generate(num_memories=num_facts)
                AgentBuilderSession.store_generated_life_facts(life_facts)
                st.rerun()
