import streamlit as st
import requests

st.set_page_config(
        page_title="IOGPT",
        page_icon="🤖",
        menu_items={}  # This helps hide the menu
    )

# Hide Streamlit menu and footer
hide_menu_style = """
        <style>
        #MainMenu {visibility: hidden;}
        footer {visibility: hidden;}
        </style>
        """
st.markdown(hide_menu_style, unsafe_allow_html=True)

class VietnameseChatbot:
    def __init__(self):
        self.api_key = st.secrets["GROQ_API_KEY"]  # Store your API key in Huggingface Secrets
        self.api_url = "https://api.groq.com/openai/v1/chat/completions"
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

    def get_response(self, user_query):
        try:
            # Add a system message to guide the model's response
            payload = {
                "model": "gemma2-9b-it",
                "messages": [
                    {
                        "role": "system", 
                        "content": "You are an assistant designed to have helpful, ethical, non-biased, and constructive conversations. Do not engage with political, sensitive, or inappropriate topics. If asked about such subjects, politely decline to discuss them and redirect the conversation to more positive, constructive topics."
                    },
                    {"role": "user", "content": user_query}
                ]
            }
            response = requests.post(
                self.api_url, headers=self.headers, json=payload
            )
            if response.status_code == 200:
                return response.json()['choices'][0]['message']['content']
            else:
                print(f"API Error: {response.status_code}")
                print(f"Response: {response.text}")
                return "Đã xảy ra lỗi khi kết nối với API. Xin vui lòng thử lại."
        except Exception as e:
            print(f"Response generation error: {e}")
            return "Đã xảy ra lỗi. Xin vui lòng thử lại."

@st.cache_resource
def initialize_chatbot():
    return VietnameseChatbot()

def main():
    st.title("🤖 Trợ Lý AI - IOGPT")
    st.caption("Trò chuyện với chúng mình nhé!")

    # Initialize chatbot using cached initialization
    chatbot = initialize_chatbot()
    
    # Chat history in session state
    if 'messages' not in st.session_state:
        st.session_state.messages = []
    
    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    
    # User input
    if prompt := st.chat_input("Hãy nói gì đó..."):
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # Display user message
        with st.chat_message("user"):
            st.markdown(prompt)
        
        # Get chatbot response
        response = chatbot.get_response(prompt)
        
        # Display chatbot response
        with st.chat_message("assistant"):
            st.markdown(response)
        
        # Add assistant message to chat history
        st.session_state.messages.append({"role": "assistant", "content": response})

if __name__ == "__main__":
    main()