import os
from typing import Dict
from datetime import datetime
import pytz

from smolagents.tools import Tool

class WeatherApiTool(Tool):
    name = "weather_api"
    description = "Fetches the weather for a given location and date."
    inputs = {'location': {'type': 'string', 'description': 'The location for which to fetch the weather.'}, 'date': {'type': 'string', 'description': 'The date for which to fetch the weather.'}}
    output_type = "any"
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.api_key = os.environ.get("OPEN_WEATHER_MAP_API_KEY")
        
        # TODO: get the timezone from the location
        self.local_tz = pytz.timezone("Europe/Brussels")

    def forward(self, location: str, date: str) -> Dict:
        import requests
        import json
        
        # Get the latitude and longitude of the location
        geocoding_url = f"http://api.openweathermap.org/geo/1.0/direct?q={location}&appid={self.api_key}"
        try:
            response = requests.get(geocoding_url)
        except requests.exceptions.RequestException as e:
            return f"Error fetching geocoding data: {str(e)}"
        geocoding_data = response.json()
        
        # TODO: match if multiple using IP country info
        lat = geocoding_data[0]["lat"]
        lon = geocoding_data[0]["lon"]
        country = geocoding_data[0]["country"]
        state = geocoding_data[0]["state"]
        
        # Get the weather data for the location
        weather_url = f"http://api.openweathermap.org/data/3.0/onecall?lat={lat}&lon={lon}&exclude=current,minutely,hourly,alerts&appid={self.api_key}"
        try:
            response = requests.get(weather_url)
        except requests.exceptions.RequestException as e:
            return f"Error fetching weather data: {str(e)}"
        weather_data = response.json()
        
        weather_data["location"] = location
        weather_data["country"] = country
        weather_data["state"] = state
        
        # Extract the weather data for the specified date
        for day in weather_data["daily"]:
            date_time = datetime.fromtimestamp(day["dt"])
            if date_time.strftime('%Y-%m-%d') == date:
                break
        
        return day