from smolagents import Tool
from typing import Any, Optional

class SimpleTool(Tool):
    name = "calculate_journey_metrics"
    description = "Calculates journey metrics including distance and time."
    inputs = {"start_location":{"type":"string","description":"The starting point of the journey."},"destination_location":{"type":"string","description":"The endpoint of the journey."},"transportation_mode":{"type":"string","nullable":True,"description":"Mode of transport ('driving', 'walking', 'bicycling', 'transit', or 'plane'). Defaults to 'driving'."}}
    output_type = "string"

    def forward(self, start_location: str, destination_location: str, transportation_mode: Optional[str] = None) -> str:
        """Calculates journey metrics including distance and time.

        Args:
            start_location: The starting point of the journey.
            destination_location: The endpoint of the journey.
            transportation_mode: Mode of transport ('driving', 'walking', 'bicycling', 'transit', or 'plane'). 
                               Defaults to 'driving'.

        Returns:
            str: Journey distance and estimated time.
        """
        from geopy.geocoders import Nominatim
        from geopy.distance import geodesic

        try:
            geolocator = Nominatim(user_agent="journey_metrics_calculator")
            start = geolocator.geocode(start_location)
            end = geolocator.geocode(destination_location)

            if not start or not end:
                return "Could not find one or both locations"

            distance_km = geodesic(
                (start.latitude, start.longitude),
                (end.latitude, end.longitude)
            ).kilometers

            speeds = {
                'driving': {'urban': 30, 'highway': 100, 'average': 80},
                'walking': {'average': 5},
                'bicycling': {'casual': 12, 'average': 20},
                'transit': {'urban': 30, 'intercity': 60},
                'plane': {
                    'short_haul': 500,
                    'medium_haul': 700,
                    'long_haul': 800
                }
            }

            if transportation_mode == 'plane':
                if distance_km < 500:
                    speed = speeds['plane']['short_haul']
                elif distance_km < 3000:
                    speed = speeds['plane']['medium_haul']
                else:
                    speed = speeds['plane']['long_haul']
                additional_time = 2
            else:
                mode = transportation_mode or 'driving'
                speed = speeds.get(mode, speeds['driving'])['average']
                additional_time = 0

            travel_time = distance_km / speed
            total_hours = travel_time + additional_time

            if total_hours < 1:
                time_str = f"{int(total_hours * 60)} minutes"
            elif total_hours < 24:
                hrs = int(total_hours)
                mins = int((total_hours - hrs) * 60)
                time_str = f"{hrs} hours and {mins} minutes"
            else:
                days = total_hours / 24
                time_str = f"{days:.1f} days"

            if transportation_mode == 'plane':
                return (
                    f"Journey metrics:\n"
                    f"Distance: {distance_km:.1f} kilometers\n"
                    f"Flight time: {int(travel_time * 60)} minutes\n"
                    f"Total estimated time (including airport procedures): {time_str}"
                )

            return (
                f"Journey metrics:\n"
                f"Distance: {distance_km:.1f} kilometers\n"
                f"Estimated time by {transportation_mode or 'driving'}: {time_str}"
            )

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