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

# ============================================
# CUSTOM TOOLS FOR ALFRED
# ============================================

@tool
def calculator(operation: str) -> str:
    """
    Performs basic math calculations.
    
    Args:
        operation: A math expression like "2 + 2" or "10 * 5"
    
    Returns:
        The result of the calculation as a string
    """
    try:
        allowed_chars = set('0123456789+-*/(). ')
        if all(c in allowed_chars for c in operation):
            result = eval(operation)
            return f"The result of {operation} is: {result}"
        else:
            return "Invalid operation. Please use only numbers and basic operators."
    except Exception as e:
        return f"Could not calculate. Error: {str(e)}"


@tool
def get_current_time() -> str:
    """
    Gets the current date and time in UTC.
    
    Returns:
        The current date and time as a string
    """
    from datetime import datetime
    now = datetime.utcnow()
    return f"Current date and time (UTC): {now.strftime('%Y-%m-%d %H:%M:%S')}"


@tool
def get_weather_info(city: str) -> str:
    """
    Gets a simple weather description for a city (mock data for demo).
    
    Args:
        city: The name of the city to get weather for
    
    Returns:
        A weather description string
    """
    # This is mock data - in production you'd use a real weather API
    import random
    conditions = ["sunny", "partly cloudy", "cloudy", "rainy", "windy"]
    temp = random.randint(15, 30)
    condition = random.choice(conditions)
    return f"Weather in {city}: {temp}°C, {condition}"


# ============================================
# MODEL SETUP - Using free-tier compatible model
# ============================================

model = InferenceClientModel(
    model_id="Qwen/Qwen2.5-7B-Instruct",  # Smaller model that works on free tier
    token=os.environ.get("HF_TOKEN"),
)

# ============================================
# ALFRED AGENT SETUP
# ============================================

alfred = CodeAgent(
    tools=[
        DuckDuckGoSearchTool(),
        calculator,
        get_current_time,
        get_weather_info,
    ],
    model=model,
    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 "Please ask me something. How may I assist you today?"
    
    try:
        response = alfred.run(message)
        return str(response)
    except Exception as e:
        error_msg = str(e)
        # Provide helpful error messages
        if "404" in error_msg or "not found" in error_msg.lower():
            return "I'm having trouble connecting to my AI backend. Please check the HF_TOKEN and model availability."
        elif "rate" in error_msg.lower() or "limit" in error_msg.lower():
            return "I've hit a rate limit. Please wait a moment and try again."
        else:
            return f"I encountered an issue: {error_msg}"


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

with gr.Blocks(title="Alfred - AI Butler Agent") as demo:
    
    gr.Markdown(
        """
        # 🎩 Alfred - Your AI Butler
        
        Good day! I am Alfred, your personal AI assistant. I am at your service.
        
        ### My Capabilities:
        - 🔍 **Web Search** - I can search the internet for information
        - 🧮 **Calculator** - I can perform mathematical calculations
        - 🕐 **Current Time** - I can tell you the current time (UTC)
        - 🌤️ **Weather** - I can provide weather information
        
        ---
        """
    )
    
    chatbot = gr.ChatInterface(
        fn=chat_with_alfred,
        examples=[
            "What is artificial intelligence?",
            "Calculate 25 * 4 + 10",
            "What time is it now?",
            "What's the weather in London?",
            "Search for the latest news about AI",
            "Calculate (100 + 50) * 2 / 5",
        ],
    )
    
    gr.Markdown(
        """
        ---
        *Built with [smolagents](https://github.com/huggingface/smolagents) and Gradio | Unit 1 Project*
        """
    )

# ============================================
# LAUNCH APP
# ============================================

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