from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import time
import requests
import pytz
import yaml
import random
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
import os

# Global cache for forecasts per city to avoid refetching within a certain duration (10 minutes)
cached_forecasts = {}  # Dictionary mapping city name -> { "data": str, "timestamp": int }
CACHE_DURATION_MS = 10 * 60 * 1000  # 10 minutes in milliseconds

@tool
def get_current_time_in_timezone(timezone: str = "UTC") -> str:
    """
    A tool to get the current time in a specific timezone.
    
    Args:
        timezone (str): The timezone string. E.g., 'America/New_York'
    
    Returns:
        str: The local time in the specified timezone.
    """
    try:
        tz = pytz.timezone(timezone)
        now = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"Local time in {timezone}: {now}"
    except Exception:
        return "Invalid timezone"

@tool
def get_weather_forecast(city: str = "New York", days: int = 7) -> str:
    """
    A tool to fetch a weekly weather forecast for a specified city.
    
    This function first checks for a cached forecast for the given city. If the cached
    data is older than 10 minutes, it attempts to fetch new data from OpenWeatherMap.
    
    If a valid API key is not provided or an error occurs during the API call,
    a mock forecast is generated.
    
    Args:
        city (str): The city name for which to get the weather.
        days (int): Number of days to forecast (maximum of 7).
    
    Returns:
        str: A formatted summary forecast with daily averages.
    """
    global cached_forecasts
    now_ms = int(time.time() * 1000)

    # Check for a valid cache
    if city in cached_forecasts and now_ms - cached_forecasts[city]["timestamp"] < CACHE_DURATION_MS:
        return cached_forecasts[city]["data"]

    # Use a real API key in production
    api_key = os.environ.get("OPENWEATHER_API", None)  # This should be defined/assigned earlier

    # For debugging: print the API key value (remove or secure in production)
    print("DEBUG: API key used =", api_key)

    # Validate API key before proceeding
    if not api_key:
        return "Error: Invalid or missing OpenWeather API key. Please configure a valid API key."

    url = f"http://api.openweathermap.org/data/2.5/forecast?q={city}&appid={api_key}&units=metric"
    try:
        response = requests.get(url)
        data = response.json()
        if data.get("cod") != "200":
            raise ValueError(data.get("message", "Error fetching weather data"))

        # Process and group forecast items by day (format: YYYY-MM-DD)
        forecast_by_day = {}
        for item in data.get("list", []):
            dt = datetime.datetime.utcfromtimestamp(item["dt"])
            day_str = dt.strftime("%Y-%m-%d")
            forecast_by_day.setdefault(day_str, []).append(item)

        # Build forecast for up to 'days' days
        sorted_days = sorted(forecast_by_day.keys())[:min(days, 7)]
        daily_forecasts = []
        for day_str in sorted_days:
            day_items = forecast_by_day[day_str]
            date_obj = datetime.datetime.strptime(day_str, "%Y-%m-%d")
            day_name = date_obj.strftime("%A")
            date_display = date_obj.strftime("%b %d")

            total_temp = total_humidity = total_pressure = total_wind = 0
            condition_counts = {}
            for item in day_items:
                total_temp += item["main"]["temp"]
                total_humidity += item["main"]["humidity"]
                total_pressure += item["main"]["pressure"]
                total_wind += item["wind"]["speed"]
                condition = item["weather"][0]["description"]
                condition_counts[condition] = condition_counts.get(condition, 0) + 1

            count = len(day_items)
            avg_temp = round(total_temp / count)
            avg_humidity = round(total_humidity / count)
            avg_pressure = round(total_pressure / count)
            avg_wind = round((total_wind / count) * 3.6)  # Convert m/s to km/h
            most_common_condition = max(condition_counts.items(), key=lambda x: x[1])[0]

            daily_forecasts.append(
                f"{day_name}, {date_display}: {avg_temp}°C, {most_common_condition} | Humidity: {avg_humidity}% | Pressure: {avg_pressure} hPa | Wind: {avg_wind} km/h"
            )
        forecast_data = "\n".join(daily_forecasts)
    except Exception as error:
        # Log the error and return an error message
        print("Error fetching forecast from API:", error)
        return f"Error: Unable to fetch weather data. Reason: {error}"

    # Cache the forecast data
    cached_forecasts[city] = {
        "data": forecast_data,
        "timestamp": now_ms
    }
    return forecast_data

final_answer = FinalAnswerTool()

model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

# Import a text-to-image tool from the hub, if desired.
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, get_current_time_in_timezone, get_weather_forecast],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()