sci/_base_class.py

# _base_class.py
import os
import re
import streamlit as st
from _arango import ArangoDB
from _chromadb import ChromaDB

class BaseClass:
    def __init__(self, username: str, **kwargs) -> None:
        self.username: str = username
        self.project_name: str = kwargs.get('project_name', None)
        self.collection: str = kwargs.get('collection_name', None)
        self.user_arango: ArangoDB = self.get_arango()
        self.base_arango: ArangoDB = self.get_arango(admin=True)
        for key, value in kwargs.items():
            setattr(self, key, value)


    def get_arango(self, admin: bool = False, db_name: str = None) -> ArangoDB:
        if db_name:
            return ArangoDB(db_name=db_name)
        elif admin:
            return ArangoDB()
        else:
            from colorprinter.print_color import print_yellow
            return ArangoDB(user=self.username, db_name=self.username)

    def get_article_collections(self) -> list:
        article_collections = self.user_arango.db.aql.execute(
            'FOR doc IN article_collections RETURN doc["name"]'
        )
        return list(article_collections)

    def get_projects(self) -> list:
        projects = self.user_arango.db.aql.execute(
            'FOR doc IN projects RETURN doc["name"]'
        )
        return list(projects)


    def get_chromadb(self):
        return ChromaDB()

    def get_project(self, project_name: str):
        doc = self.user_arango.db.aql.execute(
            f'FOR doc IN projects FILTER doc["name"] == "{project_name}" RETURN doc',
            count=True,
        )
        if doc:
            return doc.next()


    def set_filename(self, filename=None, folder="other_documents"):
        """
        Checks if a file with the given filename already exists in the download folder.
        If it does, appends or increments a numerical suffix to the filename until a unique name is found.

        Args:
            filename (str): The base name of the file to check.

        Sets:
            self.file_path (str): The unique file path generated.
        """

        download_folder = f"user_data/{self.username}/{self.document_type}"
        if self.is_sci and not self.document_type == "notes":
            self.file_path = f"sci_articles/{self.doi}.pdf".replace("/", "_")
            return os.path.exists(self.file_path)
        else:
            file_path = f"{download_folder}/{filename}"
            while os.path.exists(file_path + ".pdf"):
                if not re.search(r"(_\d+)$", file_path):
                    file_path += "_1"
                else:
                    file_path = re.sub(
                        r"(\d+)$", lambda x: str(int(x.group()) + 1), file_path
                    )
            self.file_path = file_path + ".pdf"
        return file_path


class StreamlitBaseClass(BaseClass):
    """
    StreamlitBaseClass is a base class for Streamlit applications that provides methods for managing user settings, session state, and user interactions with collections and projects.
    Methods:
        __init__(username: str, **kwargs) -> None:
            Initializes the StreamlitBaseClass with a username and additional keyword arguments.
        get_settings(field: str = None):
            Retrieves user settings from the database. If a specific field is provided, returns the value of that field. Otherwise, returns all settings.
        update_settings(key, value) -> None:
            Updates a specific setting in the database and the Streamlit session state.
        get_settings():
            Retrieves user settings from the database.
        update_session_state(page_name=None):
            Updates the Streamlit session state with the attributes of the current instance. If a page name is provided, updates the session state for that page.
        update_current_page(page_name):
            Updates the current page in the Streamlit session state and the database.
        choose_collection(text="Select a collection of favorite articles") -> str:
            Displays a select box for choosing a collection of favorite articles. Updates the current collection in the session state and the database.
        choose_project(text="Select a project") -> str:
            Displays a select box for choosing a project. Updates the current project in the session state and the database.
        """
    def __init__(self, username: str, **kwargs) -> None:
        super().__init__(username, **kwargs)
    
    def get_settings(self, field: str = None):
        """
        Retrieve or initialize user settings from the database.

        This method fetches the user settings document from the "settings" collection
        in the ArangoDB database. If the settings document does not exist, it initializes
        it with default values for "current_collection" and "current_page". The settings
        are then stored in the Streamlit session state.

        Args:
            field (str, optional): The specific field to retrieve from the settings. 
                                   If not provided, the entire settings document is returned.

        Returns:
            dict or any: The entire settings document if no field is specified, 
                         otherwise the value of the specified field.
        """
        settings = self.user_arango.db.document("settings/settings")
        if not settings:
            self.user_arango.db.collection("settings").insert(
                {"_key": "settings", "current_collection": None, "current_page": None}
            )
        for i in ["current_collection", "current_page"]:
            if i not in settings:
                settings[i] = None
        st.session_state["settings"] = settings
        if field:
            return settings[field]
        return settings

    def update_settings(self, key, value) -> None:
        """
        Update a specific setting in the database and session state.

        Args:
            key (str): The key of the setting to update.
            value (Any): The new value for the setting.

        Returns:
            None
        """
        self.user_arango.db.collection("settings").update_match(
            filters={"_key": "settings"},
            body={key: value},
            merge=True,
        )
        st.session_state["settings"][key] = value

    def update_session_state(self, page_name=None):
        """
        Updates the Streamlit session state with the attributes of the current instance.

        Parameters:
        page_name (str, optional): The name of the page to update in the session state.
                                   If not provided, it defaults to the current page stored in the session state.

        The method iterates over the instance's attributes and updates the session state
        for the given page name with those attributes that are of type str, int, float, list, dict, or bool.
        """

        if not page_name:
            page_name = st.session_state.get("current_page")
        for attr, value in self.__dict__.items():
            if any([isinstance(value, t) for t in [str, int, float, list, dict, bool]]):
                st.session_state[page_name][attr] = value
        # for k, v in st.session_state[page_name].items():
        #     if isinstance(v, list):
        #         print(k.upper())
        #         for j in v:
        #             print(j)
        #     else:
        #         print(k.upper(), v)

    def update_current_page(self, page_name):
        """
        Updates the current page in the session state and settings.

        Args:
            page_name (str): The name of the page to set as the current page.

        Side Effects:
            Updates the "current_page" in the session state and settings if it is different from the current value.
        """
        if st.session_state.get("current_page") != page_name:
            st.session_state["current_page"] = page_name
            self.update_settings("current_page", page_name)

    
    def choose_collection(self, text="Select a collection of favorite articles") -> str:
        """
        Prompts the user to select a collection of favorite articles from a list.

        Args:
            text (str): The prompt text to display for the selection box. Defaults to "Select a collection of favorite articles".

        Returns:
            str: The name of the selected collection.

        Side Effects:
            - Sets the `project` attribute to None.
            - Sets the `collection` attribute to the selected collection.
            - Updates the settings with the key "current_collection" to the selected collection.
            - Updates the session state.
        """
        collections = self.get_article_collections()
        collection = st.selectbox(text, collections, index=None)
        if collection:
            self.project = None
            self.collection = collection
            self.update_settings("current_collection", collection)
            self.update_session_state()
            return collection
    
    def choose_project(self, text="Select a project") -> str:
        """
        Prompts the user to select a project from a list of available projects.

        Args:
            text (str): The prompt text to display for project selection. Defaults to "Select a project".

        Returns:
            str: The name of the selected project.

        Side Effects:
            - Updates the current project settings.
            - Updates the session state.
            - Prints the chosen project name to the console.
        """
        projects = self.get_projects()
        print('projects', projects)
        print(self.project_name)
        
        project = st.selectbox(text, projects, index=projects.index(self.project_name) if self.project_name in projects else None)
        print('Choosing project...')
        if project:
            from projects_page import Project
            self.project = Project(self.username, project, self.user_arango)
            self.collection = None
            self.update_settings("current_project", self.project.name)
            self.update_session_state()
            print('CHOOSEN PROJECT:', self.project.name)
            return self.project