import streamlit as st
import time
from langchain.schema import HumanMessage, SystemMessage, AIMessage
from langchain.chat_models import ChatOpenAI


def get_chatmodel_response(question):
    # Retry logic
    max_retries = 3
    retries = 0

    while retries < max_retries:
        try:
            st.session_state['flowmessages'].append(HumanMessage(content=question))
            answer = chat(st.session_state['flowmessages'])
            st.session_state['flowmessages'].append(AIMessage(content=answer.content))
            return answer.content
        except Exception as e:
            print(f"Error: {e}")
            if "Rate limit" in str(e):
                print(f"Rate limit exceeded. Waiting and retrying...")
                time.sleep(5)  # Adjust the waiting time as needed
                retries += 1
            else:
                print("Unhandled exception. Please try again later.")
                break

    print("Exceeded the maximum number of retries. Please try again later.")
    return None

# Streamlit app setup
st.set_page_config(page_title="A - Square")
st.header("Welcome to A - Square")

from dotenv import load_dotenv
load_dotenv()
import os

# ChatOpenAI class
chat = ChatOpenAI(temperature=0.3)

if 'flowmessages' not in st.session_state:
    st.session_state['flowmessages'] = [
        SystemMessage(content=""" Your Name is "Ambicious Angel"/
                      
                                - You were developed by "Ambicious Angels"
                      
                                You are a story teller application. The name of the application is "A - Square"/
                                
                                As a Story teller application. You will perform the following task./
                                 1. Introduce yourself. And Interact with the user.
                                 2. The user will be a student or a kid or an adult. 
                                 3. When the user will enter a topic or askes about a toipc./
                                    i) First, You should understand the Input topic and know the meaning of it.
                                 
                                    ii) Then, You should tell him a good story about that topic.  
                                 4. You should tell story in a very proper way.
                                 5. You should be able to tell story based on any topic. 
                                 6. Also the User is interested in "Design thinking". If the user askes about it. You should be more presice and tell an interesting story.
                                 7. If the User Input is Inappropriate or if the user input is Obscene or a 18 plus content or offencive. you should tell,/
                                    "I'm sorry, but I cannot comply with that request. It's important to maintain a respectful and safe environment for all users. Is there something else I can assist you with?"
                                 8. If you did not understand the Input, Tell the user, "I'm sorry, I didn't quite understand that. Could you please provide more context or rephrase your question?"   
                                 9. After telling the story. Write the Author name as "Ambicious Angels". 
                                 10. Finally, Tell a "Thank you" and exit.
                      
                                    """)
    ]

# Streamlit UI
with st.form(key='my_form',clear_on_submit=True):
    st.markdown(
        """
        <style>
            .stTextInput {
                border-radius: 15px;
                padding: 12px;
                margin-top: 10px;
                margin-bottom: 10px;
                box-shadow: 2px 2px 5px #888888;
                border: 1px solid #dddddd;
                font-size: 16px;
                width: 100%;  /* Make the input box full width */
                height: 100px;  /* Set the height of the input box */
            }
        </style>
        """,
        unsafe_allow_html=True
    )

    input_question = st.text_input("Type here.", key="input")

    submit = st.form_submit_button("Ask A-Square")
    
    
# Add a "Clear Chat" button next to the "Submit" button
clear_chat_button = st.button("New Story", key="clear_button")

# If the "Clear Chat" button is clicked
if clear_chat_button:
    # Clear the entire session and chat
    st.session_state['flowmessages'] = []

# If the "Ask" button is clicked
if submit:
    response = get_chatmodel_response(input_question)

    if response is not None:
        # st.subheader("Here you go,")
        st.write(response)
        
    else:
        st.subheader("Error: Unable to get response. Please try again later.")