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:
            payload = {
                "model": "gemma2-9b-it",
                "messages": [
                    {"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()