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

# Fully functional weather forecast tool
@tool
def get_weather_forecast(city: str, country_code: str = "") -> str:
    """A tool that fetches the current weather and forecast for a specified location.
    
    Args:
        city: The name of the city to get weather for (e.g., 'London', 'New York')
        country_code: Optional two-letter country code (e.g., 'US', 'UK', 'FR')
    """
    try:
        # WeatherAPI.com API endpoint and parameters
        # Note: In a real application, you should store API keys securely
        API_KEY =  "397ed132ff5643738fb112308252602"  # Replace with your actual API key
        base_url = "http://api.weatherapi.com/v1"
        
        # Build location string based on inputs
        location = city
        if country_code:
            location = f"{city},{country_code}"
            
        # Create request parameters
        params = {
            'q': location,
            'appid': API_KEY,
            'units': 'metric'  # Use metric units (Celsius)
        }
        
        # Make API request
        response = requests.get(base_url, params=params)
        response.raise_for_status()  # Raise exception for HTTP errors
        
        # Parse response
        weather_data = response.json()
        
        # Extract relevant information
        current_temp = weather_data['main']['temp']
        feels_like = weather_data['main']['feels_like']
        humidity = weather_data['main']['humidity']
        conditions = weather_data['weather'][0]['description']
        wind_speed = weather_data['wind']['speed']
        
        # Format and return results
        result = (
            f"Current Weather in {weather_data['name']}, {weather_data['sys']['country']}:\n"
            f"• Temperature: {current_temp}°C (feels like {feels_like}°C)\n"
            f"• Conditions: {conditions.capitalize()}\n"
            f"• Humidity: {humidity}%\n"
            f"• Wind Speed: {wind_speed} m/s"
        )
        
        return result
    
    except requests.exceptions.HTTPError as http_err:
        return f"HTTP Error occurred: {http_err}"
    except requests.exceptions.ConnectionError:
        return f"Connection Error: Unable to connect to weather service"
    except requests.exceptions.Timeout:
        return f"Timeout Error: Request timed out"
    except requests.exceptions.RequestException as err:
        return f"Error fetching weather data: {err}"
    except KeyError as key_err:
        return f"Error parsing weather data: {key_err}"
    except Exception as e:
        return f"Unexpected 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:
        # 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()

# Use the HF endpoint as suggested in the comment since there's an auth error with the direct model
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',  # Using the suggested endpoint
    custom_role_conversions=None,
)

# Import tool from Hub
try:
    image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
except Exception as e:
    print(f"Warning: Could not load image generation tool: {e}")
    image_generation_tool = None

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

# Create a list of tools, only adding ones that are successfully loaded
tools = [final_answer, get_weather_forecast, get_current_time_in_timezone]
if image_generation_tool is not None:
    tools.append(image_generation_tool)

agent = CodeAgent(
    model=model,
    tools=tools,
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

# Modified to handle potential errors in the Gradio UI
try:
    GradioUI(agent).launch()
except TypeError as e:
    if "unsupported operand type(s) for +=" in str(e):
        print("Warning: There seems to be an issue with token counting in the Gradio UI.")
        print("This could be due to the model not properly returning token count information.")
        print("Consider modifying the Gradio_UI.py file to handle None values for token counts.")
    else:
        raise