from smolagents import CodeAgent, DuckDuckGoSearchTool, tool
from smolagents.models import LiteLLMModel  # ✨ ADD THIS
import gradio as gr
import datetime
import pytz

# ==================== YOUR TOOLS ====================
@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Get current time in any timezone.
    Args:
        timezone: Timezone like 'America/New_York' or 'Asia/Tokyo'
    """
    try:
        tz = pytz.timezone(timezone)
        time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"⏰ Time in {timezone}: {time}"
    except:
        return "Invalid timezone. Try: America/New_York, Europe/London, Asia/Tokyo"

@tool
def calculate_tip(bill_amount: float, tip_percent: float) -> str:
    """Calculate restaurant tip.
    Args:
        bill_amount: Bill total in dollars
        tip_percent: Tip percentage (e.g., 15 for 15%)
    """
    tip = bill_amount * (tip_percent / 100)
    total = bill_amount + tip
    return f"💰 Tip: ${tip:.2f}\n💵 Total: ${total:.2f}\n📊 {tip_percent}% of ${bill_amount}"

@tool
def word_counter(text: str) -> str:
    """Count words and characters.
    Args:
        text: Text to analyze
    """
    words = len(text.split())
    chars = len(text)
    return f"📝 Words: {words}\n🔤 Characters: {chars}"

@tool
def simple_calculator(expression: str) -> str:
    """Do basic math calculations.
    Args:
        expression: Math like '2 + 2' or '10 * 5 / 2'
    """
    try:
        # Safe calculation
        allowed = {'+', '-', '*', '/', '(', ')', '.', ' '}
        if all(c.isdigit() or c in allowed for c in expression):
            result = eval(expression)
            return f"🧮 {expression} = {result}"
        return "❌ Only basic math allowed: + - * / ( )"
    except:
        return "❌ Invalid math expression"

# ==================== CREATE MODEL ====================
# ✨ ADD THIS: Create a model object
model = LiteLLMModel(
    model="gpt-3.5-turbo",  # Free model that works
    api_base="https://api.openai.com/v1",  # OpenAI endpoint
    temperature=0.5
)

# ==================== CREATE AGENT ====================
agent = CodeAgent(
    model=model,  # ✨ ADD THIS: Pass the model object
    tools=[
        DuckDuckGoSearchTool(),          # Web search
        get_current_time_in_timezone,    # Time tool
        calculate_tip,                   # Tip calculator
        word_counter,                    # Word counter
        simple_calculator                # Math calculator
    ],
    max_steps=8,                         # Increased for complex tasks
    verbosity_level=1                    # Show thinking process
)

# ==================== GRADIO INTERFACE ====================
def chat_with_agent(message, history):
    """Handle chat messages"""
    response = agent.run(message)
    return response

# Create chat interface
demo = gr.ChatInterface(
    fn=chat_with_agent,
    title="🤖 Super Assistant Pro",
    description="I can: 🔍 Search web • ⏰ Tell time • 💰 Calculate tips • 📝 Count words • 🧮 Do math",
    examples=[
        "What time is it in Tokyo and calculate 15% tip on ¥5000?",
        "Count words in 'Hello world' and search for current news",
        "Calculate 25 * 4 + 10 and tell time in London",
        "What's the weather in Paris? Also calculate 18% tip on €75"
    ],
    theme="soft"
)

# Launch the app
if __name__ == "__main__":
    demo.launch()