from smolagents import CodeAgent, DuckDuckGoSearchTool, load_tool, tool
try:
    from smolagents import InferenceClientModel as ModelClass
except ImportError:
    try:
        from smolagents import HfApiModel as ModelClass
    except ImportError:
        from smolagents import HfEngine as ModelClass

import datetime
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Weather search tool using DuckDuckGo
@tool  
def get_weather_info(city: str) -> str:
    """Searches for current weather information for a given city using DuckDuckGo.
    
    Args:
        city: Name of the city to get weather information for (e.g., 'Warsaw', 'New York')
    
    Returns:
        A string containing weather information for the specified city
    """
    try:
        from duckduckgo_search import DDGS
        
        
        with DDGS() as ddgs:
            
            results = list(ddgs.text(f"{city} current weather temperature", max_results=3))
            
            if results:
                weather_info = f"Weather information for {city}:\n"
                for i, result in enumerate(results):
                    title = result.get('title', '')
                    body = result.get('body', '')
                    weather_info += f"\n{i+1}. {title}\n   {body}\n"
                return weather_info
            else:
                return f"No weather information found for {city}"
                
    except ImportError:
        # Fallback do DuckDuckGoSearchTool
        search_tool = DuckDuckGoSearchTool()
        result = search_tool(f"{city} weather")
        return f"Weather in {city}: {result}" if result else f"No weather data for {city}"
    except Exception as e:
        return f"Error: {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:
        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)}"

# Initialize tools
final_answer = FinalAnswerTool()

# Initialize model
model = ModelClass(
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
)

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

# Handle prompt templates
try:
    # Try to load custom prompts
    with open("prompts.yaml", 'r') as stream:
        prompt_templates = yaml.safe_load(stream)
    
    # Ensure all required keys are present
    if hasattr(CodeAgent, 'DEFAULT_PROMPTS'):
        for key in CodeAgent.DEFAULT_PROMPTS:
            if key not in prompt_templates:
                prompt_templates[key] = CodeAgent.DEFAULT_PROMPTS[key]
    
    # Create agent with custom prompts
    agent = CodeAgent(
        model=model,
        tools=[final_answer, get_weather_info, get_current_time_in_timezone, image_generation_tool],
        max_steps=6,
        verbosity_level=1,
        prompt_templates=prompt_templates
    )
except Exception as e:
    print(f"Error loading custom prompts: {e}")
    print("Using default prompts instead")
    
    # Create agent without custom prompts
    agent = CodeAgent(
        model=model,
        tools=[get_weather_info, get_current_time_in_timezone, image_generation_tool, final_answer],
        max_steps=6,
        verbosity_level=1
    )

# Launch Gradio UI
GradioUI(agent).launch()