import streamlit as st
from langchain.llms import HuggingFaceHub  
from langchain.memory import ConversationBufferMemory
import os
from datetime import datetime, timedelta
from dotenv import load_dotenv

# Print version information for debugging
try:
    import langchain_community

    print(f"langchain_community version: {langchain_community.__version__}")
except Exception as e:
    print(f"Error getting version info: {e}")

# Load environment variables - MUST be at the top
load_dotenv('.env')  # Explicitly load from .env file

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

# Title with better styling
st.markdown("""
    <h1 style='text-align: center; color: #2E86C1;'>
        Tourism Assistant - مساعد السياحة
    </h1>
""", unsafe_allow_html=True)

# Initialize session states
if "memory" not in st.session_state:
    st.session_state.memory = ConversationBufferMemory(
        return_messages=True,
        memory_key="chat_history"
    )

if "messages" not in st.session_state:
    st.session_state.messages = []

if "last_request" not in st.session_state:
    st.session_state.last_request = datetime.now() - timedelta(seconds=10)

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

# Get token from environment
token = os.environ.get("HUGGINGFACEHUB_API_TOKEN")

# For development, allow manual token entry if not found in environment
if not token:
    st.warning("API token not found in environment variables.")
    token = st.text_input("Enter your Hugging Face API token:", type="password")
    if not token:
        st.error("Token is required to proceed.")
        st.stop()

# Enhanced model configuration
model_config = {
    "English": {
        "repo_id": "google/flan-t5-base",  # Using a smaller model for testing
        "params": {
            "temperature": 0.7,
            "max_length": 512
        }
    },
    "العربية": {
        "repo_id": "aubmindlab/aragpt2-medium",  # Using a smaller model for testing
        "params": {
            "temperature": 0.6,
            "max_length": 1024
        }
    }
}

# Initialize the language model with enhanced error handling
try:
    # Simple initialization - let the library handle the token from env vars
    llm = HuggingFaceHub(
        repo_id=model_config[language]["repo_id"],
        model_kwargs=model_config[language]["params"]
    )

    conversation = ConversationChain(
        llm=llm,
        memory=st.session_state.memory,
        verbose=False
    )

    st.success("Model connected successfully!")

except Exception as e:
    error_msg = str(e)
    st.error(f"Error initializing model: {error_msg}")

    if "token" in error_msg.lower() or "api" in error_msg.lower():
        st.info("""
        Make sure the token is set correctly:
        1. Create a Hugging Face account and get a token
        2. Set it as an environment variable: `setx HUGGINGFACEHUB_API_TOKEN "your_token"`
        3. Restart your terminal after setting the environment variable
        """)
    st.stop()

# Display chat history with improved formatting
for message in st.session_state.messages:
    avatar = "🧑" if message["role"] == "user" else "🌍"
    with st.chat_message(message["role"], avatar=avatar):
        if language == "العربية":
            st.markdown(f"<div style='text-align: right;'>{message['content']}</div>", unsafe_allow_html=True)
        else:
            st.markdown(message["content"])

# Enhanced rate limiting
if (datetime.now() - st.session_state.last_request).seconds < 3:
    st.warning("Please wait a moment before sending another message." if language == "English"
               else "الرجاء الانتظار قليلاً قبل إرسال رسالة أخرى")
    st.stop()

# User input with better placeholder handling
prompt_placeholder = {
    "English": "Ask about destinations, culture, or safety tips...",
    "العربية": "اسأل عن الوجهات، الثقافة، أو نصائح السلامة..."
}
prompt = st.chat_input(prompt_placeholder[language])

if prompt:
    st.session_state.last_request = datetime.now()
    st.session_state.messages.append({"role": "user", "content": prompt})

    with st.chat_message("user", avatar="🧑"):
        st.markdown(prompt)

    with st.chat_message("assistant", avatar="🌍"):
        with st.spinner("Generating response..." if language == "English" else "جارٍ تحضير الرد..."):
            try:
                # Enhanced prompt engineering
                if language == "English":
                    full_prompt = """You are an expert tourism assistant specializing in:
                    - Detailed travel destination information
                    - Cultural norms and etiquette
                    - Safety recommendations
                    - Local transportation options
                    - Authentic dining experiences

                    Question: {prompt}
                    Answer in clear, detailed points:""".format(prompt=prompt)
                else:
                    full_prompt = """أنت مساعد سياحي خبير متخصص في:
                    - معلومات مفصلة عن الوجهات السياحية
                    - الأعراف الثقافية وآداب السلوك
                    - توصيات السلامة
                    - خيارات النقل المحلي
                    - تجارب تناول الطعام الأصيلة

                    السؤال: {prompt}
                    الجواب بنقاط واضحة ومفصلة:""".format(prompt=prompt)

                response = conversation.predict(input=full_prompt)

                # Post-process response
                cleaned_response = response.strip()
                if language == "العربية":
                    cleaned_response = f"<div style='text-align: right;'>{cleaned_response}</div>"

                st.markdown(cleaned_response, unsafe_allow_html=True)
                st.session_state.messages.append({"role": "assistant", "content": cleaned_response})

            except Exception as e:
                error_response = {
                    "English": f"Sorry, I encountered an error: {str(e)}. Please try again later.",
                    "العربية": f"عذرًا، حدث خطأ: {str(e)}. يرجى المحاولة مرة أخرى لاحقًا."
                }
                st.error(error_response[language])
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": error_response[language]
                })

# Sidebar with deployment-ready information
with st.sidebar:
    st.header("ℹ️ " + ("About" if language == "English" else "حول"))

    about_text = {
        "English": """
        **Tourism Expert Chatbot**  
        • Provides detailed travel information  
        • Offers cultural insights  
        • Available in English/Arabic  
        • Remembers conversation history  
        """,
        "العربية": """
        **مساعد سياحي خبير**  
        • يقدم معلومات سفر مفصلة  
        • يوفر رؤى ثقافية  
        • متاح باللغتين الإنجليزية والعربية  
        • يحفظ تاريخ المحادثة  
        """
    }

    st.markdown(about_text[language])

    # Add reset button to clear conversation
    if st.button("Reset Conversation" if language == "English" else "إعادة ضبط المحادثة"):
        st.session_state.messages = []
        st.session_state.memory.clear()
        st.experimental_rerun()

    # Add a manual token entry option in sidebar for easy updating
    st.subheader("API Configuration")
    new_token = st.text_input("Update API Token:", type="password")
    if st.button("Save Token"):
        os.environ["HUGGINGFACEHUB_API_TOKEN"] = new_token
        st.success("Token updated! Restart the app to apply.")