import streamlit as st
import streamlit_authenticator as stauth
import yaml
from yaml.loader import SafeLoader
from streamlit_authenticator import LoginError
from time import sleep

from colorprinter.print_color import *
from _arango import ArangoDB


def get_settings():
    """
    Function to get the settings from the ArangoDB using the new API.
    """
    if "username" not in st.session_state:
        return {}
        
    # Create ArangoDB instance with user's database
    arango = ArangoDB(db_name=st.session_state["username"])
    
    # Use the get_settings method from the new API
    settings = arango.get_settings()
    
    if settings:
        st.session_state["settings"] = settings
    else:
        # Initialize default settings if none exist
        default_settings = {'current_collection': None, 'current_page': None}
        arango.initialize_settings(default_settings)
        st.session_state["settings"] = default_settings
        
    return st.session_state["settings"]


st.set_page_config(page_title="Science Assistant Fish", page_icon="🐟")

with open("streamlit_users.yaml") as file:
    config = yaml.load(file, Loader=SafeLoader)

authenticator = stauth.Authenticate(
    config["credentials"],
    config["cookie"]["name"],
    config["cookie"]["key"],
    config["cookie"]["expiry_days"],
)

try:
    authenticator.login()
except LoginError as e:
    st.error(e)

if st.session_state["authentication_status"]:
    # Set username in session state
    st.session_state["username"] = st.session_state["username"]
    
    sleep(0.1)
    # Retry mechanism for importing get_settings
    for _ in range(3):
        try:
            get_settings()
        except Exception as e:
            sleep(0.3)
            print_red(f"Error getting settings: {e}")
            print("Retrying to get settings...")

    # Retry mechanism for importing pages
    for _ in range(3):
        try:
            from streamlit_pages import (
                Article_Collections,
                Bot_Chat,
                Projects,
                Settings,
                RSS_Feeds,
                Research,
                Search_Papers
            )
            break
        except ImportError as e:
            # Write the full error traceback
            sleep(0.3)
            print_red(e)
            print("Retrying to import pages...")

    st.session_state["settings"] = get_settings()
    if isinstance(st.session_state["settings"], dict) and "current_page" in st.session_state["settings"]:
        st.session_state["current_page"] = st.session_state["settings"]["current_page"]
    else:
        if "current_page" not in st.session_state:
            st.session_state["current_page"] = None

    if "not_downloaded" not in st.session_state:
        st.session_state["not_downloaded"] = {}

    # Pages
    bot_chat = st.Page(Bot_Chat)
    projects = st.Page(Projects)
    article_collections = st.Page(Article_Collections)
    settings = st.Page(Settings)
    rss_feeds = st.Page(RSS_Feeds)
    research = st.Page(Research)
    search_papers = st.Page(Search_Papers)

    sleep(0.1)
    pg = st.navigation([bot_chat, projects, article_collections, research, search_papers, rss_feeds, settings])
    sleep(0.1)
    pg.run()
    # try: #TODO Use this when in production
    #     pg.run()
    # except Exception as e:
    #     print_red(e)
    #     st.error("An error occurred. The site will be reloaded.")
    #     import traceback
    #     from datetime import datetime
    #     from time import sleep

    #     traceback_string = traceback.format_exc()
    #     traceback.print_exc()
    #     arango = ArangoDB(db_name="base")
    #     timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    #     print_rainbow(st.session_state.to_dict())
    #     session_state = st.session_state.to_dict()
    #     if 'bot' in session_state:
    #         del session_state['bot']
    #     arango.insert_document(
    #         collection_name="error_logs",
    #         document={
    #             "error": traceback_string,
    #             "_key": timestamp,
    #             "session_state": session_state,
    #         },
    #         overwrite=True
    #     )
    #     with st.status(":red[An error occurred. The site will be reloaded.]"):
    #         for i in range(5):
    #             sleep(1)
    #             st.write(f"Reloading in {5-i} seconds...")
    #     st.rerun()
    with st.sidebar:
        st.write("---")
        authenticator.logout()


elif st.session_state["authentication_status"] is False:
    st.error("Username/password is incorrect")
elif st.session_state["authentication_status"] is None:
    st.warning("Please enter your username and password")