"""Route Optimizer Agent - Optimizes delivery routes based on zones and available hubs."""
import uuid
from datetime import datetime, timedelta
from time import monotonic
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass
from math import radians, cos, sin, asin, sqrt
from collections import defaultdict

from core.agent import SpecializedAgent
from core.types import AgentTask, MessageType, ZoneType
from core.logger import logger

_CACHE_TTL_SECONDS = 600  # 10 minutes


@dataclass
class Waypoint:
    location_id: str
    lat: float
    lng: float
    address: str
    type: str  # pickup, delivery, hub, spoke
    order_id: Optional[str] = None
    time_window_start: Optional[datetime] = None
    time_window_end: Optional[datetime] = None


@dataclass
class Route:
    route_id: str
    waypoints: List[Waypoint]
    total_distance_km: float
    estimated_duration_minutes: float
    vehicle_id: str
    zones_traversed: List[str]
    fuel_cost: float
    efficiency_score: float


class RouteOptimizerAgent(SpecializedAgent):
    """Route Optimizer Agent - Optimizes delivery routes based on zones, traffic, and constraints."""

    def __init__(self):
        super().__init__(
            agent_id="ROUTE_OPTIMIZER",
            domain="route_optimization",
            description="Optimizes delivery routes based on zones, hubs, and constraints"
        )

        self._hubs = {
            "DL-HUB-01": (28.6139, 77.2090),
            "DL-HUB-02": (28.5355, 77.2100),
            "MU-HUB-01": (19.0760, 72.8777),
            "MU-HUB-02": (19.1650, 72.8500),
            "BL-HUB-01": (12.9716, 77.5946),
            "HY-HUB-01": (17.3850, 78.4867),
            "PU-HUB-01": (18.5204, 73.8567),
            "KL-HUB-01": (22.5726, 88.3639),
        }

        self._zones = self._init_zones()
        # Cache stores (Route, created_at_monotonic) — evicted after _CACHE_TTL_SECONDS
        self._route_cache: Dict[str, Tuple[Route, float]] = {}
        self._traffic_patterns = self._init_traffic_patterns()
        self._route_history: List[Dict] = []

    def _init_zones(self) -> Dict[str, Dict]:
        return {
            "north_delhi": {"pincode_range": ("100", "199"), "center": (28.6139, 77.2090), "hub": "DL-HUB-01", "typical_traffic": "medium"},
            "south_delhi": {"pincode_range": ("200", "299"), "center": (28.5355, 77.2100), "hub": "DL-HUB-02", "typical_traffic": "high"},
            "mumbai_west": {"pincode_range": ("400", "449"), "center": (19.0760, 72.8777), "hub": "MU-HUB-01", "typical_traffic": "high"},
            "mumbai_east": {"pincode_range": ("450", "499"), "center": (19.1650, 72.8500), "hub": "MU-HUB-02", "typical_traffic": "medium"},
            "bangalore":   {"pincode_range": ("560", "562"), "center": (12.9716, 77.5946), "hub": "BL-HUB-01", "typical_traffic": "medium"},
            "hyderabad":   {"pincode_range": ("500", "599"), "center": (17.3850, 78.4867), "hub": "HY-HUB-01", "typical_traffic": "medium"},
            "pune":        {"pincode_range": ("400", "499"), "center": (18.5204, 73.8567), "hub": "PU-HUB-01", "typical_traffic": "medium"},
            "kolkata":     {"pincode_range": ("600", "699"), "center": (22.5726, 88.3639), "hub": "KL-HUB-01", "typical_traffic": "low"},
        }

    def _init_traffic_patterns(self) -> Dict[str, Dict]:
        return {
            "morning": {"multiplier": 1.2, "description": "7AM-10AM rush"},
            "midday":  {"multiplier": 1.0, "description": "10AM-4PM normal"},
            "evening": {"multiplier": 1.5, "description": "4PM-8PM rush"},
            "night":   {"multiplier": 0.8, "description": "8PM-7AM light"},
        }

    # ------------------------------------------------------------------ #
    #  Cache helpers with TTL                                              #
    # ------------------------------------------------------------------ #

    def _cache_put(self, route_id: str, route: Route):
        self._route_cache[route_id] = (route, monotonic())

    def _cache_get(self, route_id: str) -> Optional[Route]:
        entry = self._route_cache.get(route_id)
        if entry is None:
            return None
        route, ts = entry
        if monotonic() - ts >= _CACHE_TTL_SECONDS:
            del self._route_cache[route_id]
            return None
        return route

    async def _heartbeat(self):
        """Evict expired entries from route cache."""
        now = monotonic()
        expired = [rid for rid, (_, ts) in self._route_cache.items() if now - ts >= _CACHE_TTL_SECONDS]
        for rid in expired:
            del self._route_cache[rid]
        if expired:
            logger.debug(f"Route cache: evicted {len(expired)} expired entries ({len(self._route_cache)} remaining)")

    # ------------------------------------------------------------------ #
    #  Task dispatch                                                        #
    # ------------------------------------------------------------------ #

    async def handle_task(self, task: AgentTask) -> Dict[str, Any]:
        handlers = {
            "optimize_route":       self._optimize_route,
            "plan_multi_stop":      self._plan_multi_stop,
            "plan_inter_hub_route": self._plan_inter_hub_route,
            "calculate_eta":        self._calculate_eta,
            "avoid_zone":           self._avoid_zone,
            "reoptimize_route":     self._reoptimize_route,
            "get_zone_routes":      self._get_zone_routes,
            "batch_optimize":       self._batch_optimize,
        }
        handler = handlers.get(task.task_type, self._unknown_task)
        return await handler(task)

    async def _optimize_route(self, task: AgentTask) -> Dict[str, Any]:
        order_id    = task.data.get("order_id")
        pickup      = task.data.get("pickup", {})
        delivery    = task.data.get("delivery", {})
        vehicle_type = task.data.get("vehicle_type", "van")

        logger.info(f"Route Optimizer: Optimizing route for order {order_id}")

        pickup_coords   = (pickup.get("lat", 28.6139),   pickup.get("lng", 77.2090))
        delivery_coords = (delivery.get("lat", 19.0760), delivery.get("lng", 72.8777))

        direct_distance  = self._haversine_distance(pickup_coords, delivery_coords)
        optimal_path     = self._find_optimal_path(pickup_coords, delivery_coords)
        total_distance   = self._calculate_total_distance(optimal_path)
        traffic_multiplier = self._get_traffic_multiplier()
        estimated_time   = (total_distance / 30) * traffic_multiplier * 60

        route_id = f"RT-OPT-{uuid.uuid4().hex[:8].upper()}"

        waypoints = []
        for i, coords in enumerate(optimal_path):
            hub_id = self._find_nearest_hub(coords)
            waypoints.append(Waypoint(
                location_id=f"WPT-{i}",
                lat=coords[0],
                lng=coords[1],
                address=str(self._hubs.get(hub_id, ("Unknown",))[0]) if hub_id else "Route point",
                type="hub" if 0 < i < len(optimal_path) - 1 else ("pickup" if i == 0 else "delivery"),
                order_id=order_id,
            ))

        route = Route(
            route_id=route_id,
            waypoints=waypoints,
            total_distance_km=total_distance,
            estimated_duration_minutes=estimated_time,
            vehicle_id=task.data.get("vehicle_id", ""),
            zones_traversed=self._identify_zones(optimal_path),
            fuel_cost=total_distance * 3.5,
            efficiency_score=self._calculate_efficiency(total_distance, direct_distance),
        )

        self._cache_put(route_id, route)

        logger.info(
            f"Route {route_id}: {len(waypoints)} waypoints | {total_distance:.1f} km | "
            f"ETA {estimated_time:.0f} min | efficiency {route.efficiency_score:.0f}%"
        )

        return {
            "status": "optimized",
            "route_id": route_id,
            "waypoints": [{"lat": w.lat, "lng": w.lng, "type": w.type, "address": w.address} for w in waypoints],
            "total_distance_km": total_distance,
            "estimated_duration_minutes": estimated_time,
            "zones_traversed": route.zones_traversed,
            "fuel_cost": route.fuel_cost,
            "efficiency_score": route.efficiency_score,
        }

    async def _plan_multi_stop(self, task: AgentTask) -> Dict[str, Any]:
        stops      = task.data.get("stops", [])
        vehicle_id = task.data.get("vehicle_id")

        logger.info(f"Route Optimizer: Planning multi-stop route with {len(stops)} stops")

        waypoints = [
            Waypoint(
                location_id=f"STOP-{i}",
                lat=stop.get("lat"),
                lng=stop.get("lng"),
                address=stop.get("address", ""),
                type=stop.get("type", "delivery"),
                order_id=stop.get("order_id"),
            )
            for i, stop in enumerate(stops)
        ]

        optimized_order = self._nearest_neighbor_optimization(waypoints)
        total_distance  = self._calculate_route_distance(optimized_order)
        estimated_time  = (total_distance / 25) * 60

        route_id = f"RT-MULTI-{uuid.uuid4().hex[:8].upper()}"
        route = Route(
            route_id=route_id,
            waypoints=optimized_order,
            total_distance_km=total_distance,
            estimated_duration_minutes=estimated_time,
            vehicle_id=vehicle_id,
            zones_traversed=self._identify_zones([(w.lat, w.lng) for w in optimized_order]),
            fuel_cost=total_distance * 3.5,
            efficiency_score=85.0,
        )
        self._cache_put(route_id, route)

        return {
            "status": "planned",
            "route_id": route_id,
            "stop_order": [{"order": i + 1, "lat": w.lat, "lng": w.lng, "type": w.type} for i, w in enumerate(optimized_order)],
            "total_distance_km": total_distance,
            "estimated_duration_minutes": estimated_time,
        }

    async def _plan_inter_hub_route(self, task: AgentTask) -> Dict[str, Any]:
        from_hub = task.data.get("from_hub")
        to_hub   = task.data.get("to_hub")
        order_id = task.data.get("order_id")

        logger.info(f"Route Optimizer: Inter-hub route {from_hub} -> {to_hub}")

        if from_hub not in self._hubs or to_hub not in self._hubs:
            return {"status": "error", "message": "Invalid hub ID(s)"}

        from_coords = self._hubs[from_hub]
        to_coords   = self._hubs[to_hub]
        direct_distance = self._haversine_distance(from_coords, to_coords)

        intermediate_hub = None
        if direct_distance > 500:
            intermediate_hub = self._find_intermediate_hub(from_coords, to_coords)

        route_coords = (
            [from_coords, self._hubs[intermediate_hub], to_coords]
            if intermediate_hub else
            [from_coords, to_coords]
        )
        total_distance = self._calculate_total_distance(route_coords)
        estimated_time = (total_distance / 40) * 60

        route_id = f"RT-IHUB-{uuid.uuid4().hex[:8].upper()}"

        return {
            "status": "planned",
            "route_id": route_id,
            "from_hub": from_hub,
            "to_hub": to_hub,
            "intermediate_hub": intermediate_hub,
            "waypoints": [{"hub": h, "coords": self._hubs.get(h, (0, 0))} for h in [from_hub, intermediate_hub, to_hub] if h],
            "total_distance_km": total_distance,
            "estimated_duration_minutes": estimated_time,
            "estimated_hours": estimated_time / 60,
        }

    async def _calculate_eta(self, task: AgentTask) -> Dict[str, Any]:
        route_id         = task.data.get("route_id")
        current_location = task.data.get("current_location")

        cached = self._cache_get(route_id)
        if cached:
            return {
                "route_id": route_id,
                "total_eta_minutes": cached.estimated_duration_minutes,
                "remaining_distance_km": cached.total_distance_km,
                "current_eta": (datetime.now() + timedelta(minutes=cached.estimated_duration_minutes)).isoformat(),
            }

        from_coords = (current_location.get("lat", 0), current_location.get("lng", 0))
        to_coords   = task.data.get("destination", (0, 0))
        distance    = self._haversine_distance(from_coords, to_coords)
        eta_minutes = (distance / 30) * self._get_traffic_multiplier() * 60

        return {
            "distance_km": distance,
            "eta_minutes": eta_minutes,
            "current_eta": (datetime.now() + timedelta(minutes=eta_minutes)).isoformat(),
        }

    async def _avoid_zone(self, task: AgentTask) -> Dict[str, Any]:
        route_id    = task.data.get("route_id")
        avoid_zone  = task.data.get("zone")

        logger.info(f"Route Optimizer: Avoiding zone {avoid_zone}")

        if self._cache_get(route_id):
            return {
                "status": "replanned",
                "route_id": route_id,
                "avoided_zone": avoid_zone,
                "additional_distance_km": 5.0,
                "additional_time_minutes": 15,
            }

        return {"status": "error", "message": "Route not found"}

    async def _reoptimize_route(self, task: AgentTask) -> Dict[str, Any]:
        route_id  = task.data.get("route_id")
        new_stops = task.data.get("new_stops", [])

        logger.info(f"Route Optimizer: Reoptimizing route {route_id}")

        route = self._cache_get(route_id)
        if route:
            for stop in new_stops:
                route.waypoints.append(Waypoint(
                    location_id=f"NEW-{len(route.waypoints)}",
                    lat=stop.get("lat"),
                    lng=stop.get("lng"),
                    address=stop.get("address", ""),
                    type="add_delivery",
                    order_id=stop.get("order_id"),
                ))
            coords = [(w.lat, w.lng) for w in route.waypoints]
            route.total_distance_km = self._calculate_total_distance(coords)
            route.estimated_duration_minutes = (route.total_distance_km / 25) * 60
            self._cache_put(route_id, route)

            return {
                "status": "reoptimized",
                "route_id": route_id,
                "new_distance_km": route.total_distance_km,
                "new_eta_minutes": route.estimated_duration_minutes,
            }

        return {"status": "error", "message": "Route not found"}

    async def _get_zone_routes(self, task: AgentTask) -> Dict[str, Any]:
        zone = task.data.get("zone")
        now  = monotonic()

        zone_routes = []
        for route_id, (route, ts) in list(self._route_cache.items()):
            if now - ts >= _CACHE_TTL_SECONDS:
                continue
            if zone in route.zones_traversed:
                zone_routes.append({
                    "route_id": route.route_id,
                    "distance_km": route.total_distance_km,
                    "duration_minutes": route.estimated_duration_minutes,
                })

        return {"zone": zone, "total_routes": len(zone_routes), "routes": zone_routes}

    async def _batch_optimize(self, task: AgentTask) -> Dict[str, Any]:
        orders = task.data.get("orders", [])
        logger.info(f"Route Optimizer: Batch optimizing {len(orders)} orders")

        zone_groups: Dict[str, list] = defaultdict(list)
        for order in orders:
            zone = self._identify_zone_from_coords((order.get("lat", 0), order.get("lng", 0)))
            zone_groups[zone].append(order)

        results = [await self._optimize_zone_routes(zone, zone_orders) for zone, zone_orders in zone_groups.items()]

        return {
            "status": "batch_optimized",
            "zones_optimized": len(results),
            "total_orders": len(orders),
            "total_distance_km": sum(r["total_distance_km"] for r in results),
            "total_time_minutes": sum(r["estimated_time_minutes"] for r in results),
            "zone_results": results,
        }

    async def _unknown_task(self, task: AgentTask) -> Dict[str, Any]:
        return {"status": "error", "message": f"Unknown task: {task.task_type}"}

    # ------------------------------------------------------------------ #
    #  Geometry helpers                                                     #
    # ------------------------------------------------------------------ #

    def _haversine_distance(self, coord1: Tuple[float, float], coord2: Tuple[float, float]) -> float:
        lat1, lon1 = coord1
        lat2, lon2 = coord2
        lat1, lon1, lat2, lon2 = map(radians, [lat1, lon1, lat2, lon2])
        dlat = lat2 - lat1
        dlon = lon2 - lon1
        a = sin(dlat / 2) ** 2 + cos(lat1) * cos(lat2) * sin(dlon / 2) ** 2
        return 2 * asin(sqrt(a)) * 6371

    def _find_optimal_path(self, start: Tuple[float, float], end: Tuple[float, float]) -> List[Tuple[float, float]]:
        start_hub = self._find_nearest_hub(start)
        end_hub   = self._find_nearest_hub(end)
        if start_hub != end_hub:
            return [start, self._hubs[start_hub], self._hubs[end_hub], end]
        return [start, end]

    def _find_nearest_hub(self, coords: Tuple[float, float]) -> Optional[str]:
        return min(self._hubs.keys(), key=lambda h: self._haversine_distance(coords, self._hubs[h]), default=None)

    def _find_intermediate_hub(self, start: Tuple[float, float], end: Tuple[float, float]) -> Optional[str]:
        mid = ((start[0] + end[0]) / 2, (start[1] + end[1]) / 2)
        return self._find_nearest_hub(mid)

    def _calculate_total_distance(self, coords: List[Tuple[float, float]]) -> float:
        return sum(self._haversine_distance(coords[i], coords[i + 1]) for i in range(len(coords) - 1))

    def _calculate_route_distance(self, waypoints: List[Waypoint]) -> float:
        return self._calculate_total_distance([(w.lat, w.lng) for w in waypoints])

    def _get_traffic_multiplier(self) -> float:
        hour = datetime.now().hour
        if 7 <= hour < 10:
            return self._traffic_patterns["morning"]["multiplier"]
        if 10 <= hour < 16:
            return self._traffic_patterns["midday"]["multiplier"]
        if 16 <= hour < 20:
            return self._traffic_patterns["evening"]["multiplier"]
        return self._traffic_patterns["night"]["multiplier"]

    def _identify_zones(self, coords: List[Tuple[float, float]]) -> List[str]:
        return list({self._identify_zone_from_coords(c) for c in coords if self._identify_zone_from_coords(c)})

    def _identify_zone_from_coords(self, coords: Tuple[float, float]) -> str:
        return min(self._zones.keys(), key=lambda z: self._haversine_distance(coords, self._zones[z]["center"]), default="unknown")

    def _calculate_efficiency(self, actual_distance: float, direct_distance: float) -> float:
        if direct_distance == 0:
            return 100.0
        return min(100.0, (direct_distance / actual_distance) * 100)

    def _nearest_neighbor_optimization(self, waypoints: List[Waypoint]) -> List[Waypoint]:
        if not waypoints:
            return []
        unvisited = waypoints[1:]
        ordered   = [waypoints[0]]
        while unvisited:
            current = ordered[-1]
            nearest = min(unvisited, key=lambda w: self._haversine_distance((current.lat, current.lng), (w.lat, w.lng)))
            ordered.append(nearest)
            unvisited.remove(nearest)
        return ordered

    async def _optimize_zone_routes(self, zone: str, orders: List[Dict]) -> Dict[str, Any]:
        total_distance = 0.0
        total_time     = 0.0
        for i in range(0, len(orders), 5):
            batch  = orders[i:i + 5]
            coords = [(o.get("lat", 0), o.get("lng", 0)) for o in batch]
            dist   = self._calculate_total_distance(coords)
            total_distance += dist
            total_time     += (dist / 25) * 60
        return {"zone": zone, "orders_in_zone": len(orders), "total_distance_km": total_distance, "estimated_time_minutes": total_time}

    async def think(self, context: str, options: List[str] = None) -> str:
        return f"[ROUTE_OPTIMIZER reasoning]: {context}"