# app.py
# Alfred - The Gala Party Agent with Agentic RAG
# Unit 3 Use Case Project for Hugging Face Agents Course

import os
import gradio as gr
from smolagents import CodeAgent, DuckDuckGoSearchTool, InferenceClientModel

# Import our custom tools
from retriever import create_guest_retriever_tool
from tools import get_current_time, calculate_party_budget, suggest_seating_arrangement, dietary_check


# ============================================
# SETUP THE MODEL
# ============================================

model = InferenceClientModel(
    model_id="Qwen/Qwen2.5-7B-Instruct",  # Free tier compatible model
    token=os.environ.get("HF_TOKEN"),
)


# ============================================
# CREATE THE GUEST RETRIEVER TOOL (RAG)
# ============================================

print("Loading guest database...")
guest_retriever_tool = create_guest_retriever_tool()
print("Guest database loaded successfully!")


# ============================================
# CREATE ALFRED - THE GALA AGENT
# ============================================

alfred = CodeAgent(
    model=model,
    tools=[
        guest_retriever_tool,      # RAG tool for guest information
        DuckDuckGoSearchTool(),    # Web search for general queries
        get_current_time,          # Time tool
        calculate_party_budget,    # Budget calculator
        suggest_seating_arrangement,  # Seating planner
        dietary_check,             # Dietary requirements helper
    ],
    max_steps=5,
    verbosity_level=1,
)


# ============================================
# CHAT FUNCTION
# ============================================

def chat_with_alfred(message, history):
    """
    Process user message and return Alfred's response.
    """
    if not message.strip():
        return "Good day! I am Alfred, your gala assistant. How may I help you with the party preparations?"
    
    try:
        # Add context to help Alfred understand his role
        enhanced_message = f"""You are Alfred, the butler and gala party assistant. 
You have access to a guest database that you can search using the guest_info_retriever tool.
You can also help with party planning, budgets, seating, and dietary requirements.

User's request: {message}

If the question is about a guest or guests, use the guest_info_retriever tool to find information.
Be helpful, polite, and thorough in your responses."""
        
        response = alfred.run(enhanced_message)
        return str(response)
    
    except Exception as e:
        error_msg = str(e)
        if "404" in error_msg or "not found" in error_msg.lower():
            return "I apologize, but I'm having trouble connecting to my systems. Please ensure the HF_TOKEN is properly configured."
        elif "rate" in error_msg.lower() or "limit" in error_msg.lower():
            return "I've reached my service limits momentarily. Please wait a moment and try again."
        else:
            return f"I encountered an unexpected issue: {error_msg}"


# ============================================
# GRADIO INTERFACE
# ============================================

with gr.Blocks(title="Alfred - Gala Party Agent") as demo:
    
    gr.Markdown(
        """
        # 🎩 Alfred - Your Gala Party Assistant
        
        Good evening! I am Alfred, your dedicated assistant for the grand gala. 
        I have access to the complete guest list and can help you with all party preparations.
        
        ### What I Can Do:
        - 🔍 **Guest Information** - Search and retrieve details about any guest
        - 🌐 **Web Search** - Find information from the internet
        - 🕐 **Current Time** - Tell you the current date and time
        - 💰 **Budget Calculator** - Calculate party expenses
        - 🪑 **Seating Planner** - Suggest seating arrangements
        - 🥗 **Dietary Helper** - Get menu suggestions for dietary requirements
        
        ---
        """
    )
    
    chatbot = gr.ChatInterface(
        fn=chat_with_alfred,
        examples=[
            "Who is Lady Galadriel and what are her interests?",
            "Which guests are vegetarian?",
            "Tell me about guests who like music",
            "Calculate budget for 50 guests at $100 each, $5000 venue, $2000 entertainment",
            "What's the current time?",
            "Suggest seating for: John, Mary, Bob, Alice, Tom, Jane at tables of 3",
            "What are vegan menu options?",
        ],
    )
    
    gr.Markdown(
        """
        ---
        *🎓 Unit 3 Agentic RAG Use Case | Hugging Face Agents Course*
        
        This agent demonstrates **Agentic RAG** - combining an AI agent with document retrieval 
        to answer questions about specific data (guest information).
        """
    )


# ============================================
# LAUNCH
# ============================================

if __name__ == "__main__":
    demo.launch()