import os
import urllib.request
import zipfile
from pathlib import Path
import pandas as pd
import requests
import yaml

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

def download_and_unzip(download_folder) -> pd.DataFrame:
    os.makedirs(download_folder, exist_ok=True)
    csv_path = Path(download_folder) / "worldcities.csv"
    if csv_path.exists():
        return pd.read_csv(csv_path, index_col=False)
    zip_filename = "simplemaps_worldcities_basicv1.77.zip"
    zip_filepath = os.path.join(download_folder, zip_filename)
    if not os.path.exists(zip_filepath):
        url = "https://simplemaps.com/static/data/world-cities/basic/simplemaps_worldcities_basicv1.77.zip"
        urllib.request.urlretrieve(url, zip_filepath)
    with zipfile.ZipFile(zip_filepath, 'r') as zip_ref:
        zip_ref.extractall(download_folder)
    return pd.read_csv(csv_path, index_col=False)

@tool
def get_weather_for_city(city: str, country: str) -> str:
    """Searches for the city in the world cities dataset, obtains its coordinates,
    and fetches current weather from the Open-Meteo API.
    Args:
        city: Name of the city (case insensitive, matches 'city_ascii' in the CSV).
        country: Name of the country (case insensitive).
    Returns:
        A string containing the city's coordinates and current weather details.
    """
    df = download_and_unzip("downloads")
    matches = df[(df['city_ascii'].str.lower() == city.lower()) & (df['country'].str.lower() == country.lower())]
    if matches.empty:
        return f"No location found for {city}, {country}."
    
    row = matches.iloc[0]
    lat = float(row['lat'])
    lng = float(row['lng'])
    
    url = "https://api.open-meteo.com/v1/forecast"
    params = {
        "latitude": lat,
        "longitude": lng,
        "current_weather": "true"
    }
    
    try:
        response = requests.get(url, params=params, timeout=10)
        if response.status_code != 200:
            return f"Error fetching weather data: {response.status_code}."
        data = response.json()
        if "current_weather" not in data:
            return "Weather data not available."
        current = data["current_weather"]
        temperature = current.get("temperature")
        windspeed = current.get("windspeed")
        winddirection = current.get("winddirection")
        time_str = current.get("time")
        return (f"Current weather in {row['city_ascii']}, {row['country']} (lat: {lat}, lng: {lng}):\n"
                f"Temperature: {temperature}°C, Wind Speed: {windspeed} km/h, Wind Direction: {winddirection}° at {time_str}.")
    except Exception as e:
        return f"An error occurred while fetching weather data: {str(e)}"

final_answer = FinalAnswerTool()

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

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

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

GradioUI(agent).launch()