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

# 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"
)

# Hugging Face API token configuration - PRODUCTION READY
HF_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
if not HF_TOKEN:
    st.error("""
        API token not configured properly. Please:
        1. Create a .env file with HUGGINGFACEHUB_API_TOKEN=your_token_here
        2. On Hugging Face Spaces, add it as a secret named HUGGINGFACEHUB_API_TOKEN
    """)
    st.stop()

# Enhanced model configuration
model_config = {
    "English": {
        "repo_id": "google/flan-t5-xxl",
        "params": {
            "temperature": 0.7,
            "max_length": 512,
            "max_new_tokens": 300,
            "repetition_penalty": 1.2
        }
    },
    "العربية": {
        "repo_id": "aubmindlab/aragpt2-mega",
        "params": {
            "temperature": 0.6,
            "max_length": 1024,
            "max_new_tokens": 400,
            "repetition_penalty": 1.3
        }
    }
}

# Initialize the language model with enhanced error handling
try:
    llm = HuggingFaceHub(
        repo_id=model_config[language]["repo_id"],
        huggingfacehub_api_token=HF_TOKEN,  # Correct parameter name
        model_kwargs=model_config[language]["params"]
    )

    conversation = ConversationChain(
        llm=llm,
        memory=st.session_state.memory,
        verbose=False
    )
except Exception as e:
    error_msg = str(e)
    if "API token" in error_msg:
        st.error("Invalid API token. Please check your Hugging Face token.")
    elif "repo_id" in error_msg:
        st.error("Model loading failed. The specified model may not be available.")
    else:
        st.error(f"Initialization error: {error_msg}")
    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": "Sorry, I encountered an error. Please try again later.",
                    "العربية": "عذرًا، حدث خطأ. يرجى المحاولة مرة أخرى لاحقًا."
                }
                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])