"""
Routing optimizer for the distribution service using Google OR-Tools VRP
"""

import logging
from typing import List, Dict, Any, Optional, Tuple
from datetime import datetime, timedelta
import time

# Google OR-Tools - Vehicle Routing Problem
try:
    from ortools.constraint_solver import routing_enums_pb2
    from ortools.constraint_solver import pywrapcp
    HAS_ORTOOLS = True
except ImportError:
    print("Warning: OR-Tools not installed. Using fallback routing algorithm.")
    HAS_ORTOOLS = False

logger = logging.getLogger(__name__)


class RoutingOptimizer:
    """
    Vehicle Routing Problem optimizer using Google OR-Tools
    """

    def __init__(self):
        self.has_ortools = HAS_ORTOOLS

    async def optimize_daily_routes(
        self, 
        deliveries: List[Dict[str, Any]], 
        depot_location: Tuple[float, float],
        vehicle_capacity_kg: Optional[float] = 1000.0,
        time_limit_seconds: float = 30.0
    ) -> Dict[str, Any]:
        """
        Optimize daily delivery routes using VRP
        
        Args:
            deliveries: List of delivery dictionaries with keys:
                - id: str - delivery ID
                - location: Tuple[float, float] - (lat, lng)
                - weight_kg: float - weight of delivery
                - time_window: Optional[Tuple[str, str]] - delivery time window
            depot_location: Tuple[float, float] - depot location (lat, lng)
            vehicle_capacity_kg: Maximum weight capacity per vehicle
            time_limit_seconds: Time limit for optimization (timeout)

        Returns:
            Dict with optimized route sequences and metadata
        """
        if not self.has_ortools:
            logger.warning("OR-Tools not available, using fallback sequential routing")
            return self._fallback_sequential_routing(deliveries, depot_location)

        start_time = time.time()

        try:
            # Prepare data for VRP
            locations = [depot_location]  # Depot is first location (index 0)
            demands = [0]  # Depot has no demand
            time_windows = [(0, 24*60)]  # Depot available all day (in minutes from midnight)

            delivery_mapping = {}
            for i, delivery in enumerate(deliveries, 1):
                locations.append(delivery['location'])
                # Ensure demands are integers for OR-Tools compatibility
                weight_kg = delivery.get('weight_kg', 0)
                demands.append(int(weight_kg) if isinstance(weight_kg, (int, float)) else 0)

                # Convert time windows to minutes from midnight
                time_window = delivery.get('time_window', None)
                if time_window:
                    start_time_str, end_time_str = time_window
                    start_minutes = self._time_to_minutes(start_time_str)
                    end_minutes = self._time_to_minutes(end_time_str)
                    time_windows.append((int(start_minutes), int(end_minutes)))
                else:
                    time_windows.append((0, 24*60))  # Default to all day if no time window

                delivery_mapping[i] = delivery['id']

            # Check if we have no deliveries (only depot), return early with empty route
            if len(locations) <= 1:  # Only depot, no deliveries
                logger.info("No deliveries to optimize, returning empty route")
                return {
                    'routes': [],
                    'total_distance_km': 0,
                    'optimization_time_seconds': time.time() - start_time,
                    'algorithm_used': 'ortools_vrp',
                    'status': 'success'
                }

            # Calculate total demand first before checking it
            total_demand = sum(demands)

            # Check if total demand is 0 but we have deliveries - handle this case too
            if total_demand == 0 and len(locations) > 1:
                logger.info("Total demand is 0 but deliveries exist, returning simple route")
                # Create simple route with all deliveries but no capacity constraints
                simple_route = {
                    'route_number': 1,
                    'route_sequence': [delivery_mapping[i] for i in range(1, len(locations))],
                    'stops': [{
                        'stop_number': i,
                        'delivery_id': delivery_mapping.get(i, f"delivery_{i}"),
                        'sequence': i - 1
                    } for i in range(1, len(locations))],
                    'total_weight_kg': 0
                }
                return {
                    'routes': [simple_route],
                    'total_distance_km': 0,
                    'optimization_time_seconds': time.time() - start_time,
                    'algorithm_used': 'ortools_vrp_zero_demand',
                    'status': 'success'
                }

            # Calculate distance matrix using haversine formula
            distance_matrix = self._calculate_distance_matrix(locations)

            # Create VRP model
            # Calculate required vehicles (total_demand already calculated above)
            # Ensure at least 1 vehicle, and enough to cover demand plus buffer
            min_vehicles = max(1, int(total_demand / vehicle_capacity_kg) + 1)
            # Add a buffer vehicle just in case
            num_vehicles = int(min_vehicles + 1)
            
            logger.info(f"VRP Optimization: Demand={total_demand}kg, Capacity={vehicle_capacity_kg}kg, Vehicles={num_vehicles}")

            # Create VRP model
            manager = pywrapcp.RoutingIndexManager(
                len(distance_matrix),  # number of locations
                num_vehicles,  # number of vehicles
                [0] * num_vehicles,  # depot index for starts
                [0] * num_vehicles   # depot index for ends
            )

            routing = pywrapcp.RoutingModel(manager)

            def distance_callback(from_index, to_index):
                """Returns the distance between the two nodes."""
                from_node = manager.IndexToNode(from_index)
                to_node = manager.IndexToNode(to_index)
                return distance_matrix[from_node][to_node]

            transit_callback_index = routing.RegisterTransitCallback(distance_callback)
            routing.SetArcCostEvaluatorOfAllVehicles(transit_callback_index)

            # Add capacity constraint
            def demand_callback(index):
                """Returns the demand of the node."""
                node = manager.IndexToNode(index)
                return int(demands[node])  # Ensure demands are integers

            demand_callback_index = routing.RegisterUnaryTransitCallback(demand_callback)
            routing.AddDimensionWithVehicleCapacity(
                demand_callback_index,
                0,  # null capacity slack
                [int(vehicle_capacity_kg)] * num_vehicles,  # vehicle maximum capacities (as integers)
                True,  # start cumul to zero
                'Capacity'
            )

            # Add time window constraint
            def time_callback(from_index, to_index):
                """Returns the travel time between the two nodes."""
                from_node = manager.IndexToNode(from_index)
                to_node = manager.IndexToNode(to_index)
                # Calculate travel time based on distance (meters) and assumed speed (km/h)
                distance_m = distance_matrix[from_node][to_node]
                distance_km = distance_m / 1000.0  # Convert meters to km
                # Assume 30 km/h average speed for city deliveries
                travel_time_minutes = (distance_km / 30.0) * 60.0
                return int(travel_time_minutes)

            time_callback_index = routing.RegisterTransitCallback(time_callback)
            routing.AddDimension(
                time_callback_index,
                60 * 24,  # Allow waiting time (24 hours in minutes)
                60 * 24,  # Maximum time per vehicle (24 hours in minutes)
                False,  # Don't force start cumul to zero
                'Time'
            )
            time_dimension = routing.GetDimensionOrDie('Time')

            # Add time window constraints for each location
            for location_idx in range(len(locations)):
                index = manager.NodeToIndex(location_idx)
                if index != -1:  # Valid index
                    min_time, max_time = time_windows[location_idx]
                    time_dimension.CumulVar(index).SetRange(int(min_time), int(max_time))

            # Setting first solution heuristic
            search_parameters = pywrapcp.DefaultRoutingSearchParameters()
            search_parameters.first_solution_strategy = (
                routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC
            )
            search_parameters.time_limit.FromSeconds(time_limit_seconds)
            
            # Solve the problem
            solution = routing.SolveWithParameters(search_parameters)
            
            # Check if solution was found
            if solution:
                optimized_routes = self._extract_routes(routing, manager, solution, delivery_mapping)
                
                # Calculate total distance and duration
                total_distance = 0
                total_duration = 0
                for route in optimized_routes:
                    route_distance = 0
                    for stop in route['stops']:
                        route_distance += stop.get('distance_to_next', 0)
                    route['total_distance_km'] = route_distance
                    total_distance += route_distance
                
                logger.info(f"VRP optimization completed in {time.time() - start_time:.2f}s")
                
                return {
                    'routes': optimized_routes,
                    'total_distance_km': total_distance,
                    'optimization_time_seconds': time.time() - start_time,
                    'algorithm_used': 'ortools_vrp',
                    'status': 'success'
                }
            else:
                logger.warning("OR-Tools failed to find solution, using fallback routing")
                return self._fallback_sequential_routing(deliveries, depot_location)
                
        except Exception as e:
            logger.error(f"Error in VRP optimization: {e}")
            # Fallback to simple sequential routing
            return self._fallback_sequential_routing(deliveries, depot_location)

    def _calculate_distance_matrix(self, locations: List[Tuple[float, float]]) -> List[List[int]]:
        """
        Calculate distance matrix using haversine formula (in meters)
        """
        import math
        
        def haversine_distance(lat1, lon1, lat2, lon2):
            """Calculate distance between two lat/lon points in meters"""
            R = 6371000  # Earth's radius in meters
            
            lat1, lon1, lat2, lon2 = map(math.radians, [lat1, lon1, lat2, lon2])
            
            dlat = lat2 - lat1
            dlon = lon2 - lon1
            
            a = math.sin(dlat/2)**2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlon/2)**2
            c = 2 * math.asin(math.sqrt(a))
            
            return R * c  # Distance in meters
        
        n = len(locations)
        matrix = [[0] * n for _ in range(n)]
        
        for i in range(n):
            for j in range(n):
                if i != j:
                    lat1, lon1 = locations[i]
                    lat2, lon2 = locations[j]
                    dist_m = haversine_distance(lat1, lon1, lat2, lon2)
                    matrix[i][j] = int(dist_m)
        
        return matrix

    def _extract_routes(self, routing, manager, solution, delivery_mapping) -> List[Dict[str, Any]]:
        """
        Extract routes from OR-Tools solution
        """
        routes = []

        for vehicle_id in range(manager.GetNumberOfVehicles()):
            index = routing.Start(vehicle_id)

            # Skip if vehicle is not used (Start -> End directly)
            if routing.IsEnd(solution.Value(routing.NextVar(index))):
                continue

            current_route = {
                'route_number': vehicle_id + 1,
                'stops': [],
                'total_weight_kg': 0
            }

            # Initialize route sequence to store the delivery IDs in visit order
            route_sequence = []

            # Add depot as first stop
            node_index = manager.IndexToNode(index)
            delivery_id = delivery_mapping.get(node_index, f"depot_{node_index}")

            current_route['stops'].append({
                'stop_number': 1,
                'delivery_id': delivery_id,
                'location': 'depot',
                'sequence': 0
            })

            stop_number = 1

            while not routing.IsEnd(index):
                index = solution.Value(routing.NextVar(index))
                node_index = manager.IndexToNode(index)

                if node_index != 0:  # Not depot
                    stop_number += 1
                    delivery_id = delivery_mapping.get(node_index, f"delivery_{node_index}")
                    current_route['stops'].append({
                        'stop_number': stop_number,
                        'delivery_id': delivery_id,
                        'location_index': node_index,
                        'sequence': stop_number
                    })

                    # Add delivery ID to route sequence (excluding depot stops)
                    route_sequence.append(delivery_id)
                else:  # Back to depot
                    stop_number += 1
                    current_route['stops'].append({
                        'stop_number': stop_number,
                        'delivery_id': f"depot_end_{vehicle_id + 1}",
                        'location': 'depot',
                        'sequence': stop_number
                    })
                    break

            # Add the route_sequence to the current route
            current_route['route_sequence'] = route_sequence
            routes.append(current_route)

        return routes

    def _time_to_minutes(self, time_str: str) -> int:
        """
        Convert HH:MM string to minutes from midnight
        """
        if ":" in time_str:
            hour, minute = map(int, time_str.split(":"))
            return hour * 60 + minute
        else:
            # If it's already in minutes, return as is
            return int(time_str)

    def _fallback_sequential_routing(self, deliveries: List[Dict[str, Any]], depot_location: Tuple[float, float]) -> Dict[str, Any]:
        """
        Fallback routing algorithm that sequences deliveries sequentially
        """
        import math
        
        def haversine_distance(lat1, lon1, lat2, lon2):
            """Calculate distance between two lat/lon points in km"""
            R = 6371  # Earth's radius in km
            
            lat1, lon1, lat2, lon2 = map(math.radians, [lat1, lon1, lat2, lon2])
            
            dlat = lat2 - lat1
            dlon = lon2 - lon1
            
            a = math.sin(dlat/2)**2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlon/2)**2
            c = 2 * math.asin(math.sqrt(a))
            
            return R * c  # Distance in km

        # Calculate distances from depot to each delivery and between deliveries
        deliveries_with_distance = []
        for delivery in deliveries:
            lat, lon = delivery['location']
            depot_lat, depot_lon = depot_location
            dist = haversine_distance(depot_lat, depot_lon, lat, lon)
            deliveries_with_distance.append({
                **delivery,
                'distance_from_depot': dist
            })
        
        # Sort deliveries by distance from depot (nearest first)
        deliveries_with_distance.sort(key=lambda x: x['distance_from_depot'])
        
        # Create simple route
        route_stops = []
        total_distance = 0
        
        # Start from depot
        route_stops.append({
            'stop_number': 1,
            'delivery_id': 'depot_start',
            'location': depot_location,
            'sequence': 0,
            'is_depot': True
        })
        
        # Add deliveries
        for i, delivery in enumerate(deliveries_with_distance, 1):
            route_stops.append({
                'stop_number': i + 1,
                'delivery_id': delivery['id'],
                'location': delivery['location'],
                'weight_kg': delivery.get('weight_kg', 0),
                'sequence': i,
                'is_depot': False
            })
        
        # Return to depot
        route_stops.append({
            'stop_number': len(deliveries_with_distance) + 2,
            'delivery_id': 'depot_end',
            'location': depot_location,
            'sequence': len(deliveries_with_distance) + 1,
            'is_depot': True
        })
        
        # Calculate total distance
        for i in range(len(route_stops) - 1):
            current_stop = route_stops[i]
            next_stop = route_stops[i + 1]
            
            if not current_stop['is_depot'] or not next_stop['is_depot']:
                if not current_stop['is_depot'] and not next_stop['is_depot']:
                    # Between two deliveries
                    curr_lat, curr_lon = current_stop['location']
                    next_lat, next_lon = next_stop['location']
                    dist = haversine_distance(curr_lat, curr_lon, next_lat, next_lon)
                elif current_stop['is_depot'] and not next_stop['is_depot']:
                    # From depot to delivery
                    depot_lat, depot_lon = current_stop['location']
                    del_lat, del_lon = next_stop['location']
                    dist = haversine_distance(depot_lat, depot_lon, del_lat, del_lon)
                elif not current_stop['is_depot'] and next_stop['is_depot']:
                    # From delivery to depot
                    del_lat, del_lon = current_stop['location']
                    depot_lat, depot_lon = next_stop['location']
                    dist = haversine_distance(del_lat, del_lon, depot_lat, depot_lon)
                else:
                    dist = 0  # depot to depot
                    
                total_distance += dist
                route_stops[i]['distance_to_next'] = dist

        # Create route sequence from delivery IDs in the order they appear
        route_sequence = [stop['delivery_id'] for stop in route_stops if not stop.get('is_depot', False)]

        return {
            'routes': [{
                'route_number': 1,
                'stops': route_stops,
                'route_sequence': route_sequence,
                'total_distance_km': total_distance,
                'total_weight_kg': sum(d.get('weight_kg', 0) for d in deliveries),
            }],
            'total_distance_km': total_distance,
            'optimization_time_seconds': 0,
            'algorithm_used': 'fallback_sequential',
            'status': 'success'
        }