import gradio as gr
from smolagents import GradioUI, CodeAgent, HfApiModel
import ast

# Import our custom tools from their modules
from tools import DuckDuckGoSearchTool, WeatherInfoTool, HubStatsTool
from retriever import load_guest_dataset

# Initialize the Hugging Face model
model = HfApiModel()

# Initialize the web search tool
search_tool = DuckDuckGoSearchTool()

# Initialize the weather tool
weather_info_tool = WeatherInfoTool()

# Initialize the Hub stats tool
hub_stats_tool = HubStatsTool()

# Load the guest dataset and initialize the guest info tool
guest_tool = load_guest_dataset()

# Make guest tool description VERY clear
guest_tool.description = """USE THIS TOOL FOR GALA GUEST QUESTIONS!
Official guest database with names, relations, descriptions, emails.
Query examples: "all guests", "Lady Ada Lovelace", "scientists".
Returns formatted guest information."""

# Simple agent that formats output
class SimpleAgent(CodeAgent):
    def run(self, query: str):
        response = super().run(query)
        
        # Handle list responses
        if isinstance(response, list):
            if response:
                formatted = "🎩 **Gala Guests:**\n\n"
                for guest in response:
                    formatted += f"• {guest}\n"
                return formatted
            return "No guests found."
        
        # Handle string responses that contain lists
        if isinstance(response, str):
            # Check if it's a string representation of a list
            response = response.strip()
            if (response.startswith('[') and response.endswith(']')) or \
               (response.startswith("['") and response.endswith("']")):
                try:
                    guest_list = ast.literal_eval(response)
                    if isinstance(guest_list, list):
                        formatted = "🎩 **Gala Guests:**\n\n"
                        for guest in guest_list:
                            formatted += f"• {guest}\n"
                        return formatted
                except:
                    pass
            
            # Clean up code execution traces
            lines = response.split('\n')
            clean_lines = []
            for line in lines:
                if not any(line.startswith(x) for x in ['━━━━', '─ Executing', '```', 'Execution logs:']):
                    clean_lines.append(line)
            return '\n'.join(clean_lines).strip()
        
        return str(response)

# Create Alfred
alfred = SimpleAgent(
    tools=[guest_tool, weather_info_tool, hub_stats_tool, search_tool], 
    model=model,
    add_base_tools=True,
    planning_interval=2,
    max_steps=3
)

# Test
print("\n🎩 Testing Alfred...")
test_response = alfred.run("Who is coming to the gala?")
print(f"Test response:\n{test_response}")

# Create Gradio interface
with gr.Blocks(title="🎩 Alfred Assistant") as demo:
    gr.Markdown("# 🎩 Alfred Assistant")
    gr.Markdown("Ask about gala guests, weather, model stats, or search the web.")
    
    chatbot = gr.Chatbot(height=400)
    msg = gr.Textbox(label="Your question...")
    clear = gr.ClearButton([msg, chatbot])
    
    def respond(message, chat_history):
        response = alfred.run(message)
        chat_history.append((message, response))
        return "", chat_history
    
    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    
    # Examples
    with gr.Accordion("📋 Example Questions", open=True):
        gr.Examples(
            examples=[
                "Who is coming to the gala?",
                "Tell me about Lady Ada Lovelace",
                "What's the weather in London?",
                "Show me stats for GPT-2",
                "Search for Python programming news"
            ],
            inputs=msg,
            label="Click any example:"
        )

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