import streamlit as st
import os

from dotenv import load_dotenv
from rag import Rag
from vectore_store.PineconeConnector import PineconeConnector
from vectore_store.VectoreStoreManager import VectoreStoreManager
from subscription.index import get_usages, is_usage_limit_reached, daily_limit_enabled

from util import getYamlConfig

load_dotenv()

GROUP_NAME = os.environ.get("APP_NAME")
LOGO = "assets/logo.png"

COMPANIES = []

def init_app():

    config = getYamlConfig()

    if len(st.session_state) == 0:
        # Define Vectore store strategy
        pinecone_connector = PineconeConnector()
        vs_manager = VectoreStoreManager(pinecone_connector)

        st.session_state["retrived_documents"] = []
        st.session_state["messages"] = []
        st.session_state["remove_undefined_value"] = True
        st.session_state["assistant"] = Rag(vectore_store=vs_manager)
        st.session_state["data_dict"] = {}
        st.session_state["company_code"] = None
        st.session_state["prompt_system"] = config['prompt_system']
        st.session_state["chat_history_afp"] = []
        st.session_state["chat_history_cettons"] = []
        st.session_state["chat_history_steel"] = []

        # Initialize data_dict for each company
        for company in config['companies']:
            COMPANIES.append(company)
            company_code = company['code']

            st.session_state["data_dict"][company_code] = []

            if 'parts' in company['variables']:
                # Flatten structure by adding part name to each field
                st.session_state["data_dict"][company_code] = [
                    {**field, "part": part["name"]}
                    for part in company["variables"]["parts"]
                    for field in part["fields"]
                ]

            else:
                # Initialize session state with single list of variables
                st.session_state["data_dict"][company_code] = [
                    {**field} for field in company["variables"]
                ]


def main():

    init_app()

    st.set_page_config(page_title=GROUP_NAME)

    st.logo(LOGO)
    st.title(GROUP_NAME)
    
    if daily_limit_enabled():
        usages = get_usages()
        reached = is_usage_limit_reached()

        st.sidebar.html(f"<p style='margin:0;text-align:center;font-weight:bold;font-size:24px;'>{usages['count']} / {usages['daily_limit']}</p>")
        if reached:
            st.sidebar.html(f"<p style='text-align:center;font-style:italic;'>Vous avez atteint la limite d'utilisation<p>")
        else:
            st.sidebar.html(f"<p style='text-align:center;'>Nombre de messages envoyés aujourd'hui<p>")

    

    saved_documents = st.Page("pages/persistent_documents.py", title="Communs", icon="🗃️")
    documents = st.Page("pages/documents.py", title="Vos documents", icon="📂")
    prompt_system = st.Page("pages/prompt_system.py", title="Prompt système", icon="🖊️")
    # chatbot = st.Page("pages/chatbot.py", title="Chatbot", icon="🤖", default=True)

    pg = st.navigation(
        {
            "Documents": [
                saved_documents,
                documents,
            ],
            "Configurations": [
                prompt_system,
            ],
            "Auto France Parts": [
                st.Page("pages/afp.py", title="Paramètres", icon="📋"),
                st.Page("pages/chatbot_afp.py", title="Dialogue", icon="🤖", default=True),
            ],
            "La Pyramide des Cettons": [
                st.Page("pages/cettons.py", title="Paramètres", icon="📋"),
                st.Page("pages/chatbot_cettons.py", title="Dialogue", icon="🤖"),

            ],
            "Steel France": [
                st.Page("pages/steel.py", title="Paramètres", icon="📋"),
                st.Page("pages/chatbot_steel.py", title="Dialogue", icon="🤖"),

            ]
        }
    )

    pg.run()


if __name__ == "__main__":
    main()