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

@tool
def travel_search(city: str) -> str:
    """A tool that searches for travel information about a city.
    
    Args:
        city: Name of the city to search for
    """
    web_search = DuckDuckGoSearchTool()
    try:
        city_info = web_search(f"{city} tourism highlights attractions")
        return f"Travel information for {city}:\n{city_info}"
    except Exception as e:
        return f"Error searching for {city}: {str(e)}"

@tool
def airport_time_diff(airport1: str, airport2: str) -> str:
    """Calculate time difference between two airports and show attractions at destination.
    
    Args:
        airport1: First airport (city or IATA code)
        airport2: Second airport (city or IATA code)
    """
    web_search = DuckDuckGoSearchTool()
    
    def get_airport_timezone(airport):
        try:
            search_result = web_search(f"{airport} airport timezone")
            timezone_map = {
                'EST': 'America/New_York',
                'EDT': 'America/New_York',
                'CST': 'America/Chicago',
                'CDT': 'America/Chicago',
                'PST': 'America/Los_Angeles',
                'PDT': 'America/Los_Angeles',
                'CET': 'Europe/Berlin',
                'CEST': 'Europe/Berlin',
                'GMT': 'Europe/London',
                'BST': 'Europe/London',
                'JST': 'Asia/Tokyo'
            }
            
            for tz_code, tz_name in timezone_map.items():
                if tz_code in search_result.upper():
                    return tz_name
            return 'UTC'
            
        except Exception as e:
            return 'UTC'

    # Get timezones and current times
    tz1 = get_airport_timezone(airport1)
    tz2 = get_airport_timezone(airport2)
    time1 = datetime.datetime.now(pytz.timezone(tz1))
    time2 = datetime.datetime.now(pytz.timezone(tz2))
    
    # Get destination city attractions
    try:
        # Extract city name from airport2 (remove 'airport' if present)
        dest_city = airport2.lower().replace('airport', '').strip()
        attractions = web_search(f"{dest_city} top tourist attractions highlights")
    except Exception as e:
        attractions = "Could not fetch attractions information."

    return f"""
Time Comparison:
🛫 {airport1}: {time1.strftime('%H:%M')} ({tz1})
🛬 {airport2}: {time2.strftime('%H:%M')} ({tz2})

Popular Attractions at {airport2}:
{attractions}
"""

@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:
        tz = pytz.timezone(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()

model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
    custom_role_conversions=None
)

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, 
        DuckDuckGoSearchTool(),
        get_current_time_in_timezone,
        travel_search,
        airport_time_diff,
        image_generation_tool
    ],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

if __name__ == "__main__":
    interface = GradioUI(agent)
    interface.launch()