from typing import Dict, List

import streamlit as st

from synthetic_agents.model.entity.life_memory import LifeMemory


class AgentBuilderSession:
    """
    This class handles web session data storage and retrieval pertained to the agent builder page.
    """

    _USER_PROFILE_KEY = "agent_attributes"
    _LIFE_FACTS_KEY = "life_facts"
    _EDITION_ON_KEY = "edition_on"

    @staticmethod
    def initialize():
        """
        Creates an agent builder session and initializes the streamlit session object with the
        required keys.
        """
        keys = [
            AgentBuilderSession._USER_PROFILE_KEY,
            AgentBuilderSession._LIFE_FACTS_KEY,
        ]

        for key in keys:
            if key not in st.session_state:
                st.session_state[key] = None

        if AgentBuilderSession._EDITION_ON_KEY not in st.session_state:
            st.session_state[AgentBuilderSession._EDITION_ON_KEY] = False

    @staticmethod
    def get_stored_user_profile() -> Dict[str, str]:
        """
        Gets the user profile stored in the session.

        :return: stored generated user profile.
        """
        return st.session_state[AgentBuilderSession._USER_PROFILE_KEY]

    @staticmethod
    def get_stored_life_facts() -> List[LifeMemory]:
        """
        Gets the life facts stored in the session.

        :return: stored life facts.
        """
        return st.session_state[AgentBuilderSession._LIFE_FACTS_KEY]

    @staticmethod
    def is_edition_on() -> Dict[str, str]:
        """
        Gets a boolean indicating whether the agent is in edition mode.

        :return boolean indicating whether the agent is in edition mode.
        """
        return st.session_state[AgentBuilderSession._EDITION_ON_KEY]

    @staticmethod
    def store_user_profile(user_profile: Dict[str, str]):
        """
        Stores user profile in the session.

        :param user_profile: agent attributes to be stored.
        """

        st.session_state[AgentBuilderSession._USER_PROFILE_KEY] = user_profile

    @staticmethod
    def store_generated_life_facts(life_facts: List[LifeMemory]):
        """
        Stores life facts in the session.

        :param life_facts: life facts to be stored.
        """

        st.session_state[AgentBuilderSession._LIFE_FACTS_KEY] = life_facts

    @staticmethod
    def toggle_agent_edition(edition_on: bool):
        """
        Stores a boolean indicating whether the agent is in edition mode.

        :param edition_on: whether edition mode is on.
        """
        st.session_state[AgentBuilderSession._EDITION_ON_KEY] = edition_on
