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


# ── Tool 1: Timezone checker ──────────────────────────────────────────────────
@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:
        tz = pytz.timezone(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)}"


# ── Tool 2: Weather checker ───────────────────────────────────────────────────
@tool
def get_weather(city: str) -> str:
    """Fetches the current weather for a given city using the Open-Meteo API.
    Args:
        city: The name of the city to get weather for (e.g., 'London').
    """
    try:
        # Step 1: Geocode the city name to lat/lon
        geo_url = f"https://geocoding-api.open-meteo.com/v1/search?name={city}&count=1"
        geo_resp = requests.get(geo_url, timeout=10).json()
        if not geo_resp.get("results"):
            return f"Could not find location for city: {city}"
        loc = geo_resp["results"][0]
        lat, lon, name = loc["latitude"], loc["longitude"], loc["name"]

        # Step 2: Fetch weather
        weather_url = (
            f"https://api.open-meteo.com/v1/forecast"
            f"?latitude={lat}&longitude={lon}"
            f"¤t_weather=true&hourly=relative_humidity_2m"
        )
        weather_resp = requests.get(weather_url, timeout=10).json()
        cw = weather_resp.get("current_weather", {})
        temp   = cw.get("temperature", "N/A")
        wind   = cw.get("windspeed", "N/A")
        code   = cw.get("weathercode", "N/A")
        return (
            f"Weather in {name}:\n"
            f"  Temperature : {temp} °C\n"
            f"  Wind speed  : {wind} km/h\n"
            f"  Weather code: {code} (WMO standard)"
        )
    except Exception as e:
        return f"Error fetching weather: {str(e)}"


# ── Tool 3: Unit converter ────────────────────────────────────────────────────
@tool
def convert_units(value: float, from_unit: str, to_unit: str) -> str:
    """Converts a value between common units (length, weight, temperature).
    Args:
        value: The numeric value to convert.
        from_unit: The source unit (e.g., 'km', 'kg', 'celsius', 'miles', 'pounds', 'fahrenheit').
        to_unit: The target unit (e.g., 'miles', 'pounds', 'fahrenheit').
    """
    conversions = {
        ("km", "miles"): lambda x: x * 0.621371,
        ("miles", "km"): lambda x: x * 1.60934,
        ("kg", "pounds"): lambda x: x * 2.20462,
        ("pounds", "kg"): lambda x: x * 0.453592,
        ("meters", "feet"): lambda x: x * 3.28084,
        ("feet", "meters"): lambda x: x * 0.3048,
        ("celsius", "fahrenheit"): lambda x: x * 9 / 5 + 32,
        ("fahrenheit", "celsius"): lambda x: (x - 32) * 5 / 9,
        ("celsius", "kelvin"): lambda x: x + 273.15,
        ("kelvin", "celsius"): lambda x: x - 273.15,
        ("liters", "gallons"): lambda x: x * 0.264172,
        ("gallons", "liters"): lambda x: x * 3.78541,
    }
    key = (from_unit.lower(), to_unit.lower())
    if key in conversions:
        result = conversions[key](value)
        return f"{value} {from_unit} = {round(result, 4)} {to_unit}"
    return f"Conversion from '{from_unit}' to '{to_unit}' is not supported."


# ── Tool 4: Day-of-week calculator ───────────────────────────────────────────
@tool
def get_day_of_week(date_str: str) -> str:
    """Returns the day of the week for a given date.
    Args:
        date_str: A date string in YYYY-MM-DD format (e.g., '2025-12-25').
    """
    try:
        date_obj = datetime.datetime.strptime(date_str, "%Y-%m-%d")
        day_name = date_obj.strftime("%A")
        return f"{date_str} falls on a {day_name}."
    except ValueError:
        return "Invalid date format. Please use YYYY-MM-DD (e.g., '2025-07-04')."


# ── Tool 5: Random joke fetcher ───────────────────────────────────────────────
@tool
def get_random_joke(category: str) -> str:
    """Fetches a random joke from the JokeAPI.
    Args:
        category: Joke category — one of 'Programming', 'Misc', 'Dark', 'Pun', 'Spooky', 'Christmas'.
    """
    try:
        url = f"https://v2.jokeapi.dev/joke/{category}?blacklistFlags=nsfw,racist,sexist"
        resp = requests.get(url, timeout=10).json()
        if resp.get("type") == "single":
            return resp["joke"]
        elif resp.get("type") == "twopart":
            return f"{resp['setup']}\n... {resp['delivery']}"
        return "Couldn't fetch a joke right now."
    except Exception as e:
        return f"Error fetching joke: {str(e)}"


# ── Model & agent setup ───────────────────────────────────────────────────────
final_answer = FinalAnswerTool()

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

# Load image generation tool from the Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

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

agent = CodeAgent(
    model=model,
    tools=[
        final_answer,
        DuckDuckGoSearchTool(),       # Web search
        image_generation_tool,        # Text-to-image
        get_current_time_in_timezone, # Timezone lookup
        get_weather,                  # Live weather
        convert_units,                # Unit converter
        get_day_of_week,              # Date → weekday
        get_random_joke,              # Joke fetcher
    ],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates,
)

GradioUI(agent).launch()