#    _____ _    _ _____    _____ _           _   _           _   
#   / ____| |  | |_   _|  / ____| |         | | | |         | |  
#  | |  __| |  | | | |   | |    | |__   __ _| |_| |__   ___ | |_ 
#  | | |_ | |  | | | |   | |    | '_ \ / _` | __| '_ \ / _ \| __|
#  | |__| | |__| |_| |_  | |____| | | | (_| | |_| |_) | (_) | |_ 
#   \_____|\____/|_____|  \_____|_| |_|\__,_|\__|_.__/ \___/ \__|
#                                                         
#    __________  ___  ___________  _____   ___ 
#   /  _/_  __/ / _ )/ __/_  __| \/ / _ | / _ \
#  _/ /  / /   / _  / _/  / /   \  / __ |/ , _/
# /___/ /_/   /____/___/ /_/    /_/_/ |_/_/|_| 
#
# SOFT PILLOW VERZIÓ
# HF - Streamlit verzió

import streamlit as st
import time
from openai import OpenAI
import os

# Get secrets from environment variables
api_key = os.environ.get('OPENAI_API_KEY')
assistant_id = os.environ.get('ASSISTANT_ID')

# Initialize OpenAI client
client = OpenAI(api_key=api_key)

# Verify credentials are loaded
if not api_key or not assistant_id:
    st.error("Please set up OPENAI_API_KEY and ASSISTANT_ID in your Hugging Face Space secrets.")
    st.stop()

# Custom avatars
assistant_avatar = "https://huggingface.co/spaces/InfoBetyar/RAG-Chat/resolve/main/puhaparnaicon.webp"
user_avatar = "https://huggingface.co/spaces/InfoBetyar/RAG-Chat/resolve/main/usericon.webp" 

# Center-align the image and headers
col1, col2 = st.columns([1, 2])  # The ratio 1:2 gives the second column more space

# Left column: Image
with col1:
    st.image("https://huggingface.co/spaces/InfoBetyar/RAG-Chat/resolve/main/puhaparna.webp", width=200)

# Right column: Titles
with col2:
    st.markdown("""
        <div style="margin-top: -30px;">  <!-- Added padding to vertically align with image -->
            <h1 style="margin: 0;">Soft Pillow Apartment - Chatbot</h1>
            <h3 style="margin: 0;">IT Betyár - Minta alkalmazás diákoknak</h3>
        </div>
        """, unsafe_allow_html=True)

# Initialize messages in session state if not present
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display existing messages
for message in st.session_state.messages:
    with st.chat_message(message["role"], avatar=assistant_avatar if message["role"] == "assistant" else user_avatar):
        st.write(message["content"])

# Chat input
if prompt := st.chat_input("Type your message..."):
    # Display user message
    with st.chat_message("user", avatar=user_avatar):
        st.write(prompt)
    
    # Add user message to state
    st.session_state.messages.append({"role": "user", "content": prompt})

    # Get assistant response
    with st.chat_message("assistant", avatar=assistant_avatar):
        with st.spinner("Thinking..."):
            try:
                # Create thread and send message
                thread = client.beta.threads.create()
                client.beta.threads.messages.create(
                    thread_id=thread.id,
                    role="user",
                    content=prompt
                )

                # Run assistant
                run = client.beta.threads.runs.create(
                    thread_id=thread.id,
                    assistant_id=assistant_id
                )

                # Wait for completion
                while run.status != "completed":
                    time.sleep(0.5)
                    run = client.beta.threads.runs.retrieve(
                        thread_id=thread.id,
                        run_id=run.id
                    )

                # Get and display response
                messages = client.beta.threads.messages.list(thread_id=thread.id)
                response = messages.data[0].content[0].text.value
                cleaned_response = response.split("【")[0]
                st.write(cleaned_response)

                # Add assistant response to state
                st.session_state.messages.append(
                    {"role": "assistant", "content": cleaned_response}
                )

                # Clean up thread
                client.beta.threads.delete(thread.id)
                
            except Exception as e:
                st.error(f"An error occurred: {str(e)}")

# Teszt üzenetek:
#-------------------------------------
# What's the check in time?
# What are the check-in and check-out times?
# Are there any local restaurants you'd suggest?
# What's the Wi-Fi password for the property?
# What's the pin for the lockbox?
# What is the check-out procedure?
# What is the manager's name?     -->   NEM kell tudja!