import streamlit as st
import streamlit_authenticator as stauth
from streamlit_authenticator import Authenticate
from herbal_expert import herbal_expert
import yaml

def authentication_page():
    with open('creds.yaml') as f:
        creds = yaml.load(f, Loader=yaml.loader.SafeLoader)

    print("creds: ", creds)

    authenticator = Authenticate(
        creds['credentials'],
        creds['cookie']['name'],
        creds['cookie']['key'],
        creds['cookie']['expiry_days'],
        creds['preauthorized']
    )

    name, authentication_status, username = authenticator.login('Login', 'main')
    print("name: ", name)
    print("authentication_status: ", authentication_status)
    print("username: ", username)
    if authentication_status:
        authenticator.logout('Logout', 'main')
    if st.session_state["authentication_status"]:
        st.session_state.is_authenticated = True

def chatbot_page():
    st.title("Herbal Expert Chatbot")
    # Store LLM generated responses
    if "messages" not in st.session_state.keys():
        st.session_state.messages = [{"role": "assistant", "content": "How may I help you?"}]

    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.write(message["content"])


    # Function for generating LLM response
    def generate_response(prompt_input):
        print(st.session_state.messages)
        response = herbal_expert.query_expert(prompt_input)
        return response['response']


    # User-provided prompt
    if prompt := st.chat_input():
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.write(prompt)

    # Generate a new response if last message is not from assistant
    if st.session_state.messages[-1]["role"] != "assistant":
        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                response = generate_response(prompt)
                st.write(response)
        message = {"role": "assistant", "content": response}
        st.session_state.messages.append(message)

if __name__ == "__main__":
    st.session_state.is_authenticated = False
    authentication_page()

    # Check if the user is authenticated before displaying the chatbot page
    if st.session_state.is_authenticated:
        chatbot_page()