import gradio as gr
import time
from typing import List, Dict, Any, Optional

class GradioUI:
    def __init__(self, agent):
        self.agent = agent
        self.chat_history = []
        
    def process_agent_response(self, response):
        """Process the agent's response for display in the UI."""
        if isinstance(response, str):
            return response
        elif hasattr(response, "content"):
            return response.content
        else:
            return str(response)
    
    def chat(self, message, history):
        """Process a user message and return the agent's response."""
        self.chat_history = history
        
        # Add user message to history
        self.chat_history.append((message, ""))
        
        # Get response from agent
        response = self.agent.run(message)
        processed_response = self.process_agent_response(response)
        
        # Update the last response in history
        self.chat_history[-1] = (message, processed_response)
        
        return "", self.chat_history
    
    def launch(self, **kwargs):
        """Launch the Gradio interface."""
        with gr.Blocks(css="footer {visibility: hidden}") as demo:
            gr.Markdown("# 🤖 City Weather")
            gr.Markdown("""
            This agent can:
            - Get weather information for cities (using OpenWeatherMap API)
            - Search the web using DuckDuckGo
            - Provide natural conversational responses
            
            Weather Tool Limitations:
            - Cannot disambiguate between cities with the same name
            - Only accepts letters, digits, spaces, and hyphens in city names
            - Does not support additional location information (state, country, etc.)
            - Returns temperatures in Fahrenheit and wind speeds in mph
            
            Example queries:
            - "What's the weather like in Tokyo?"
            - "Tell me the temperature in New York"
            - "Search for the latest news about AI"
            
            Note: Make sure both HUGGINGFACE_TOKEN and OPENWEATHERMAP_API_KEY environment variables are set.
            """)
            
            chatbot = gr.Chatbot(
                [],
                elem_id="chatbot",
                avatar_images=(None, "🤖"),
                bubble_full_width=False,
                height=500
            )
            
            with gr.Row():
                txt = gr.Textbox(
                    scale=4,
                    show_label=False,
                    placeholder="Enter your message here...",
                    container=False
                )
                submit_btn = gr.Button("Send", variant="primary", scale=1)
            
            txt.submit(self.chat, [txt, chatbot], [txt, chatbot])
            submit_btn.click(self.chat, [txt, chatbot], [txt, chatbot])
            
            gr.Markdown("### 📝 Created as part of the [Hugging Face Agents Course](https://huggingface.co/learn/agents)")
            
        return demo.launch(**kwargs)