from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool 
from Gradio_UI import GradioUI

# Updated with a useful function that provides information about a location
@tool
def my_custom_tool(location: str, details_required: str = "general") -> str:
    """A tool that fetches information about a specified location.
    
    Args:
        location: The name of the location to get information about
        details_required: The type of information required (general, weather, population)
    """
    try:
        # This is a simple implementation - in a real scenario you might use an API
        if details_required.lower() == "weather":
            # Simulate a weather API call
            weather_conditions = ["sunny", "cloudy", "rainy", "snowy", "windy"]
            import random
            condition = random.choice(weather_conditions)
            temp = random.randint(0, 35)  # Celsius
            return f"Current weather in {location}: {condition}, {temp}°C"
        
        elif details_required.lower() == "population":
            # Simulate population data lookup
            # In real implementation, you would fetch from a database or API
            populations = {
                "new york": "8.4 million",
                "tokyo": "14 million",
                "london": "9 million",
                "paris": "2.1 million",
                "sydney": "5.3 million"
            }
            
            location_lower = location.lower()
            if location_lower in populations:
                return f"The population of {location} is approximately {populations[location_lower]}."
            else:
                return f"Population data for {location} not found in the database."
        
        else:  # general information
            return f"Here is general information about {location}. This is a placeholder - in a real implementation, this would fetch actual data about the location."
            
    except Exception as e:
        return f"Error fetching information for '{location}': {str(e)}"

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        # Create timezone object
        tz = pytz.timezone(timezone)
        
        # Get current time in that timezone
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"

final_answer = FinalAnswerTool()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud'
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',  # it is possible that this model may be overloaded
    custom_role_conversions=None,
)

# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

agent = CodeAgent(
    model=model,
    tools=[final_answer, my_custom_tool, get_current_time_in_timezone, image_generation_tool],  # add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()