import gradio as gr
from gradio import ChatMessage
import time
import random

def calculator_tool(expression):
    """Simple calculator tool"""
    try:
        return eval(expression)
    except:
        return "Error: Could not evaluate expression"

def weather_tool(location):
    """Simulated weather tool"""
    conditions = ["sunny", "cloudy", "rainy", "snowy"]
    temp = random.randint(0, 35)
    return f"{random.choice(conditions)}, {temp}°C in {location}"

def agent_with_tools(query, history):
    
    # Initial thinking
    thinking = ChatMessage(
        role="assistant",
        content="Let me think about this query...",
        metadata={"title": "🧠 Thinking", "id": 1}
    )
    history.append(thinking)
    yield history
    
    # Decide which tool to use based on query
    if "calculate" in query.lower() or any(op in query for op in ['+', '-', '*', '/']):
        # Extract the expression (simplified for demo)
        expression = query.split("calculate")[-1].strip() if "calculate" in query.lower() else query
        
        # Show tool usage as nested thought
        tool_call = ChatMessage(
            role="assistant",
            content=f"Expression to evaluate: {expression}",
            metadata={
                "title": "🧮 Calculator Tool", 
                "parent_id": 1,
                "id": 2,
                "status": "pending"
            }
        )
        history.append(tool_call)
        yield history
        
        # Simulate calculation time
        time.sleep(1)
        
        # Get result and update tool call
        result = calculator_tool(expression)
        tool_call.content = f"Expression: {expression}\nResult: {result}"
        tool_call.metadata["status"] = "done"
        tool_call.metadata["duration"] = 0.8  # Simulated duration
        yield history
        
        # Final response
        response = ChatMessage(
            role="assistant",
            content=f"I calculated that for you. The result is {result}."
        )
        
    elif "weather" in query.lower():
        # Extract location (simplified)
        location = query.split("weather in")[-1].strip() if "weather in" in query.lower() else "your location"
        
        # Show tool usage
        tool_call = ChatMessage(
            role="assistant",
            content=f"Checking weather for: {location}",
            metadata={
                "title": "🌤️ Weather Tool", 
                "parent_id": 1,
                "id": 2,
                "status": "pending"
            }
        )
        history.append(tool_call)
        yield history
        
        # Simulate API call
        time.sleep(1.5)
        
        # Get result and update tool call
        result = weather_tool(location)
        tool_call.content = f"Location: {location}\nCurrent conditions: {result}"
        tool_call.metadata["status"] = "done"
        tool_call.metadata["duration"] = 1.2  # Simulated duration
        yield history
        
        # Final response
        response = ChatMessage(
            role="assistant",
            content=f"I checked the weather for you. It's currently {result}."
        )
    else:
        # Default response for other queries
        time.sleep(1)
        response = ChatMessage(
            role="assistant",
            content=f"I understand you're asking about '{query}', but I don't have a specific tool for that. I can help with calculations or weather."
        )
    
    # Add final response
    history.append(response)
    yield history

demo = gr.ChatInterface(
    agent_with_tools,
    title="Sample Agents with Tool Visualization using gr.ChatMessage",
    description="Ask me to calculate something or check the weather!",
    examples=[
        "Calculate 145 * 32", 
        "What's the weather in Tokyo?",
        "Tell me about quantum physics"
    ],
    type="messages"
)

demo.launch()