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

# Fetch METAR/TAF aviation weather data
@tool
def get_aviation_weather(icao_code: str) -> str:
    """Fetches the METAR and TAF weather reports for a given airport ICAO code.
    Args:
        icao_code: The ICAO airport code (e.g., 'JFK' for John F. Kennedy International).
    """
    try:
        url = f"https://aviationweather.gov/metar/data?ids={icao_code}&format=raw&taf=on"
        response = requests.get(url)
        if response.status_code == 200:
            return response.text
        else:
            return f"Error fetching aviation weather for {icao_code}"
    except Exception as e:
        return f"Failed to retrieve data: {str(e)}"

# Flight time estimation tool
@tool
def estimate_flight_time(distance_nm: float, speed_kts: float) -> str:
    """Estimates flight time based on distance (nautical miles) and speed (knots).
    Args:
        distance_nm: Distance in nautical miles.
        speed_kts: Speed in knots.
    """
    if speed_kts <= 0:
        return "Invalid speed. Must be greater than zero."
    time_hours = distance_nm / speed_kts
    return f"Estimated flight time: {time_hours:.2f} hours."

# Fuel consumption estimator
@tool
def estimate_fuel_burn(rate_gph: float, time_hours: float) -> str:
    """Estimates total fuel consumption for a flight.
    Args:
        rate_gph: Fuel burn rate in gallons per hour.
        time_hours: Estimated flight duration in hours.
    """
    fuel_needed = rate_gph * time_hours
    return f"Estimated fuel needed: {fuel_needed:.2f} gallons."

# Current time in major aviation hubs
@tool
def get_airport_local_time(icao_code: str) -> str:
    """Gets the local time of an airport using its ICAO code.
    Args:
        icao_code: ICAO airport code (e.g., 'LHR' for London Heathrow).
    """
    airport_timezones = {
        "JFK": "America/New_York",
        "LAX": "America/Los_Angeles",
        "LHR": "Europe/London",
        "DXB": "Asia/Dubai",
        "HND": "Asia/Tokyo"
    }
    if icao_code not in airport_timezones:
        return "Airport timezone not available."
    
    tz = pytz.timezone(airport_timezones[icao_code])
    local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
    return f"Local time at {icao_code}: {local_time}"

# Finalizing agent setup
final_answer = FinalAnswerTool()
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
)

agent = CodeAgent(
    model=model,
    tools=[final_answer, get_aviation_weather, estimate_flight_time, estimate_fuel_burn, get_airport_local_time],
    max_steps=6,
    verbosity_level=1,
    name="AeroNavBot",
    description="An AI assistant for aviation navigation, weather, and flight planning.",
)

GradioUI(agent).launch()