import streamlit as st
from langchain.llms import HuggingFaceHub
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
import os

# Configure page
st.set_page_config(page_title="Tourism Chatbot", page_icon="🌍")

# Title
st.title("Tourism Assistant - مساعد السياحة")

# Initialize session state for memory
if "memory" not in st.session_state:
    st.session_state.memory = ConversationBufferMemory(return_messages=True)

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

# Language selector
language = st.selectbox("Choose Language / اختر اللغة", ["English", "العربية"])

# Hugging Face API token
hf_api_token = os.environ.get("HF_API_TOKEN", "")
if not hf_api_token:
    st.error("Hugging Face API token not found. Please set it in the Space secrets.")
    st.stop()

# Select the appropriate model based on language
model_name = "bigscience/bloom" if language == "English" else "bigscience/bloom"

# Initialize the language model
llm = HuggingFaceHub(
    repo_id=model_name,
    huggingface_api_token=hf_api_token,
    model_kwargs={"temperature": 0.7, "max_length": 512}
)

# Initialize conversation chain
conversation = ConversationChain(
    llm=llm,
    memory=st.session_state.memory,
    verbose=True
)

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

# User input
if language == "English":
    prompt = st.chat_input("Ask about tourist destinations, recommendations, or travel tips...")
else:
    prompt = st.chat_input("اسأل عن الوجهات السياحية أو التوصيات أو نصائح السفر...")

# Process the user input
if prompt:
    # 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)

    # Display assistant response with a spinner
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            # Add context for tourism assistant
            if language == "English":
                full_prompt = f"You are a helpful tourism assistant. Answer the following question about travel and tourism: {prompt}"
            else:
                full_prompt = f"أنت مساعد سياحي مفيد. أجب عن السؤال التالي حول السفر والسياحة: {prompt}"

            response = conversation.predict(input=full_prompt)
            st.markdown(response)

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

# Sidebar with information
with st.sidebar:
    st.header("About this Chatbot")
    if language == "English":
        st.write(
            "This is a tourism assistant that can help you with travel recommendations, destination information, and travel tips.")
        st.write("Feel free to ask questions in English or Arabic!")
    else:
        st.write("هذا مساعد سياحي يمكنه مساعدتك في توصيات السفر ومعلومات الوجهة ونصائح السفر.")
        st.write("لا تتردد في طرح الأسئلة باللغة الإنجليزية أو العربية!")