import datetime
import requests
import pytz
import yaml
import os
from typing import Dict, Union

from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
from tools.final_answer import FinalAnswerTool


import time



HF_TOKEN = os.getenv("HF_TOKEN", "default_hf_token_if_not_found")

@tool
def get_weather(city: str) -> str:
    """Fetches the current weather for a specified city using Open-Meteo API (No API Key Needed).
    Args:
        city: The name of the city.
    """
    try:
        # Get city coordinates from Open-Meteo's geocoding API
        geo_url = f"https://geocoding-api.open-meteo.com/v1/search?name={city}&count=1"
        geo_response = requests.get(geo_url)
        
        if geo_response.status_code == 200 and geo_response.json().get("results"):
            city_data = geo_response.json()["results"][0]
            lat, lon = city_data["latitude"], city_data["longitude"]
        else:
            return f"Error: Unable to find location for '{city}'."

        # Fetch weather data
        weather_url = f"https://api.open-meteo.com/v1/forecast?latitude={lat}&longitude={lon}&current_weather=true"
        weather_response = requests.get(weather_url)

        if weather_response.status_code == 200:
            weather_data = weather_response.json()["current_weather"]
            temp = weather_data["temperature"]
            wind_speed = weather_data["windspeed"]
            
            return f"""Weather in {city}:
            Temperature: {temp}°C
            Wind Speed: {wind_speed} m/s"""
        else:
            return "Error fetching weather data."

    except Exception as e:
        return f"Error: {str(e)}"

# Initialize tools
final_answer = FinalAnswerTool()

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

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

# Initialize the agent
agent = CodeAgent(
    model=model,
    tools=[final_answer, get_weather],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    prompt_templates=prompt_templates
)


if __name__ == "__main__":
    from Gradio_UI import demo  # Import Gradio app
    demo.launch(server_name="0.0.0.0", server_port=7860)