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
        import time

        try:
            # Initialize geocoder with increased timeout
            geolocator = Nominatim(user_agent="journey_metrics_calculator", timeout=10)

            # Add delay between requests to avoid rate limiting
            time.sleep(1)
            start = geolocator.geocode(start_location)
            time.sleep(1)
            end = geolocator.geocode(destination_location)

            if not start or not end:
                return "Could not find one or both locations. Please check spelling and try again."

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

            # Define speeds (km/h) with more realistic values
            speeds = {
                'driving': {'urban': 50, 'highway': 100, 'average': 80},
                'walking': {'average': 5},
                'bicycling': {'casual': 15, 'average': 20},
                'transit': {'urban': 30, 'intercity': 60},
                'plane': {
                    'short_haul': 500,  # <500km
                    'medium_haul': 700, # 500-3000km
                    'long_haul': 850    # >3000km
                }
            }

            # Calculate time based on mode and distance
            if transportation_mode == 'plane':
                # Select appropriate plane speed based on distance
                if distance_km < 500:
                    speed = speeds['plane']['short_haul']
                elif distance_km < 3000:
                    speed = speeds['plane']['medium_haul']
                else:
                    speed = speeds['plane']['long_haul']

                # Calculate flight time
                flight_hours = distance_km / speed

                # Add airport procedures time (2 hours total)
                total_hours = flight_hours + 2

                # Format the response for plane travel
                flight_minutes = int(flight_hours * 60)
                total_hrs = int(total_hours)
                total_mins = int((total_hours - total_hrs) * 60)

                return (
                    f"Journey metrics:\n"
                    f"Distance: {distance_km:.1f} kilometers\n"
                    f"Flight time: {flight_minutes} minutes\n"
                    f"Total estimated time (including airport procedures): "
                    f"{total_hrs} hours and {total_mins} minutes"
                )
            else:
                # Get speed for other modes
                mode = transportation_mode or 'driving'
                if mode not in speeds:
                    return f"Unsupported transportation mode: {mode}"

                speed = speeds[mode]['average']
                hours = distance_km / speed

                # Format time string
                hrs = int(hours)
                mins = int((hours - hrs) * 60)

                return (
                    f"Journey metrics:\n"
                    f"Distance: {distance_km:.1f} kilometers\n"
                    f"Estimated time by {mode}: {hrs} hours and {mins} minutes"
                )

        except Exception as e:
            return f"Error calculating journey metrics: Please try again with more specific location names. Error: {str(e)}"