from typing import Any, Optional
from smolagents.tools import Tool
import pytz
import datetime
import re
from duckduckgo_search import DDGS

class CityTimeTool(Tool):
    name = "city_time"
    description = "Get the current local time for any city by searching for its timezone and converting current UTC time."
    inputs = {'city': {'type': 'string', 'description': 'The name of the city to get the current time for'}}
    output_type = "string"
    
    def __init__(self):
        super().__init__()
        self.ddgs = DDGS()
        
    def _find_timezone(self, city: str) -> str:
        """Search for the timezone of a given city."""
        # Search for timezone information
        query = f"{city} timezone UTC time zone"
        results = list(self.ddgs.text(query, max_results=3))
        
        if not results:
            raise Exception(f"Could not find timezone information for {city}")
            
        # Common timezone patterns in city names
        common_cities = {
            'new york': 'America/New_York',
            'london': 'Europe/London',
            'paris': 'Europe/Paris',
            'tokyo': 'Asia/Tokyo',
            'sydney': 'Australia/Sydney',
            'dubai': 'Asia/Dubai',
            'singapore': 'Asia/Singapore',
            'hong kong': 'Asia/Hong_Kong',
            'los angeles': 'America/Los_Angeles',
            'chicago': 'America/Chicago',
            'toronto': 'America/Toronto',
            'sao paulo': 'America/Sao_Paulo',
            'rio': 'America/Sao_Paulo',
            'moscow': 'Europe/Moscow',
            'berlin': 'Europe/Berlin',
            'madrid': 'Europe/Madrid',
            'rome': 'Europe/Rome',
            'beijing': 'Asia/Shanghai',
            'shanghai': 'Asia/Shanghai',
            'seoul': 'Asia/Seoul',
            'bangkok': 'Asia/Bangkok',
            'cairo': 'Africa/Cairo',
            'johannesburg': 'Africa/Johannesburg',
            'mexico city': 'America/Mexico_City',
            'buenos aires': 'America/Argentina/Buenos_Aires'
        }
        
        # Check if city is in common cities first
        city_lower = city.lower()
        if city_lower in common_cities:
            return common_cities[city_lower]
            
        # Try to extract timezone from search results
        for result in results:
            text = result['body'].lower()
            # Look for common timezone patterns
            for known_city, timezone in common_cities.items():
                if known_city in text and city_lower in text:
                    return timezone
                    
        # Default to UTC if we can't determine the timezone
        raise Exception(f"Could not determine timezone for {city}. Please try a major city name.")

    def forward(self, city: str) -> str:
        try:
            # Find the timezone
            timezone_str = self._find_timezone(city)
            
            # Get current time in that timezone
            tz = pytz.timezone(timezone_str)
            utc_now = datetime.datetime.now(pytz.UTC)
            local_time = utc_now.astimezone(tz)
            
            return f"Current time in {city} ({timezone_str}) is: {local_time.strftime('%Y-%m-%d %H:%M:%S %Z')}"
            
        except Exception as e:
            return f"Error getting time for {city}: {str(e)}"