import logging
import random
import time
import math
from math import radians, cos, sin, asin, sqrt, ceil
from typing import List, Dict, Any, Optional, Set
from collections import defaultdict
from app.config.rider_preferences import RIDER_PREFERRED_KITCHENS, BLOCKED_RIDERS as _BLOCKED_RIDERS_SET
from app.services.routing.gps_smoother import (
    smooth_rider_locations,
    smooth_order_coordinates,
)
from app.config.dynamic_config import get_config

logger = logging.getLogger(__name__)


def _zone_key(lat: float, lon: float) -> str:
    """O(1) ~5 km grid cell key for zone proximity matching."""
    if lat == 0 or lon == 0:
        return "unknown"
    return f"{int(lat / 0.044)},{int(lon / 0.044)}"


class AssignmentService:
    def __init__(self):
        # Curated config drives HARD kitchen ownership. Copy so substitution
        # mutations are per-request and never bleed into the module-level dict.
        self.rider_preferences = dict(RIDER_PREFERRED_KITCHENS)
        # SOFT steering (preference bonus, distance bypass, home bonus) uses the
        # learned-augmented view: curated config ∪ learned affinity. Falls back to
        # pure config if the affinity service is unavailable.
        try:
            from app.services.routing.rider_affinity_service import get_rider_affinity
            _aff = get_rider_affinity()
            self.soft_preferences = _aff.get_preferred_kitchens()
            self.home_locations = _aff.get_home_locations()
        except Exception:
            from app.config.rider_preferences import RIDER_HOME_LOCATIONS as _H
            self.soft_preferences = dict(RIDER_PREFERRED_KITCHENS)
            self.home_locations = dict(_H)

        # Apply today's substitutions: copy absent rider's profile onto sub rider.
        # The sub rider appears in getriderlogs; the absent rider does not.
        # Reverts automatically the next day because the lookup is date-keyed.
        try:
            from app.services.rider.substitution_service import get_substitution_service
            _sub_map = get_substitution_service().get_today_map()  # {absent_id: sub_id}
            if _sub_map:
                for _absent_id, _sub_id in _sub_map.items():
                    # Hard ownership: sub rider inherits absent rider's kitchens
                    _absent_hard = self.rider_preferences.get(_absent_id, [])
                    if _absent_hard:
                        _sub_hard = list(self.rider_preferences.get(_sub_id, []))
                        _hard_lower = {k.lower() for k in _sub_hard}
                        for _k in _absent_hard:
                            if _k.lower() not in _hard_lower:
                                _sub_hard.append(_k)
                        self.rider_preferences[_sub_id] = _sub_hard
                    # Soft preferences: merge absent rider's learned kitchens
                    _absent_soft = self.soft_preferences.get(_absent_id, [])
                    if _absent_soft:
                        _sub_soft = list(self.soft_preferences.get(_sub_id, []))
                        _soft_lower = {k.lower() for k in _sub_soft}
                        for _k in _absent_soft:
                            if _k.lower() not in _soft_lower:
                                _sub_soft.append(_k)
                        self.soft_preferences[_sub_id] = _sub_soft
                    # Home location: use absent rider's home if sub has none
                    _absent_home = self.home_locations.get(_absent_id, (0.0, 0.0))
                    if _absent_home and _absent_home != (0.0, 0.0):
                        _sub_home = self.home_locations.get(_sub_id, (0.0, 0.0))
                        if not _sub_home or _sub_home == (0.0, 0.0):
                            self.home_locations[_sub_id] = _absent_home
                logger.info(
                    f"[Substitution] {len(_sub_map)} active substitution(s) applied: "
                    + ", ".join(f"absent={a} → sub={s}" for a, s in _sub_map.items())
                )
        except Exception as _e:
            logger.warning(f"[Substitution] Could not apply substitutions: {_e}")

        self.earth_radius_km = 6371
        self._cfg = get_config()

        # Cost parameters for composite cost function
        self._fuel_rate = 2.5  # Per km
        self._base_rider_cost = 0.0
        self._merchant_margin_avg = 5.0  # Default average margin

    def calculate_order_profit_features(
        self, order: Dict[str, Any], distance_km: float
    ) -> Dict[str, float]:
        """
        Calculate profit and profit-density for one order, used as a scoring
        signal (profit_bonus) when picking which rider gets a cluster.
        """
        try:
            order_amount = float(
                order.get("orderamount") or order.get("deliveryamount") or 0
            )
        except:
            order_amount = 0.0

        # Rider cost: base + (distance * fuel_rate)
        rider_cost = self._base_rider_cost + (distance_km * self._fuel_rate)

        # Profit = revenue - cost
        profit = order_amount - rider_cost

        # Profit density = profit / distance — high density means profitable
        # short deliveries; used to prioritise clusters worth serving.
        profit_density = profit / distance_km if distance_km > 0 else 0.0

        return {"profit": profit, "profit_density": profit_density}

    def _load_config(self):
        """Load ML-tuned hyperparams fresh on every assignment call."""
        cfg = self._cfg
        self.MAX_PICKUP_DISTANCE_KM = cfg.get("max_pickup_distance_km")
        self.MAX_KITCHEN_DISTANCE_KM = cfg.get("max_kitchen_distance_km")

        # Hard cap: ML hypertuner must never exceed 12 orders per rider.
        # Food delivery reality — more than 12 stops means 90+ min routes,
        # cold food, and rider exhaustion.  The tuner found 20 optimal for
        # some statistical metric but that is operationally unacceptable.
        _raw_max = int(cfg.get("max_orders_per_rider"))
        self.MAX_ORDERS_PER_RIDER = min(_raw_max, 12)

        # Ideal load must be ≥ half of max so fill-before-spill is meaningful.
        # ideal_load=4 with max=12 means riders switch at 33% — too early.
        # Floor at 6 so a rider reaches 50% before we open a new one.
        _raw_ideal = int(cfg.get("ideal_load"))
        self.IDEAL_LOAD = max(_raw_ideal, 6)

        # ── DYNAMIC NEW-RIDER PENALTY ─────────────────────────────────────
        # WORKLOAD_PENALTY_WEIGHT must be loaded FIRST so that NEW_RIDER_PENALTY
        # (which depends on it) can be computed correctly.
        #
        # Bug that was here: NEW_RIDER_PENALTY was calculated on the line ABOVE
        # the line that sets WORKLOAD_PENALTY_WEIGHT, causing AttributeError on
        # the first call and stale values on every subsequent call.
        #
        # Fix: NEW_RIDER_PENALTY = WORKLOAD_PENALTY_WEIGHT × 1.5
        # This guarantees non-preferred riders NEVER beat a preferred rider
        # at ANY load level (max preferred score = 76.82, penalty = 115.2).
        # Preferred riders fill to capacity before others are ever opened.
        self.WORKLOAD_BALANCE_THRESHOLD = cfg.get("workload_balance_threshold")
        self.WORKLOAD_PENALTY_WEIGHT = cfg.get("workload_penalty_weight")   # ← set FIRST
        self.NEW_RIDER_PENALTY = self.WORKLOAD_PENALTY_WEIGHT * 1.5         # ← then derived
        self.DISTANCE_PENALTY_WEIGHT = cfg.get("distance_penalty_weight")
        self.PREFERENCE_BONUS = cfg.get("preference_bonus")
        self.HOME_ZONE_BONUS_4KM = cfg.get("home_zone_bonus_4km")
        self.HOME_ZONE_BONUS_2KM = cfg.get("home_zone_bonus_2km")
        self.EMERGENCY_LOAD_PENALTY = cfg.get("emergency_load_penalty")

    def haversine(self, lat1, lon1, lat2, lon2):
        """Calculate the great circle distance between two points."""
        lon1, lat1, lon2, lat2 = map(
            radians, [float(lon1), float(lat1), float(lon2), float(lat2)]
        )
        dlon = lon2 - lon1
        dlat = lat2 - lat1
        a = sin(dlat / 2) ** 2 + cos(lat1) * cos(lat2) * sin(dlon / 2) ** 2
        c = 2 * asin(min(1.0, sqrt(a)))  # Clamp to 1.0 to avoid domain errors
        return c * self.earth_radius_km

    def get_lat_lon(self, obj: Dict[str, Any], prefix: str = "") -> tuple[float, float]:
        """Generic helper to extract lat/lon from diversely named keys."""
        # Try specific prefixes first
        candidates = [
            (f"{prefix}lat", f"{prefix}lon"),
            (f"{prefix}lat", f"{prefix}long"),
            (f"{prefix}latitude", f"{prefix}longitude"),
        ]
        # Also try standard keys if prefix fails
        candidates.extend(
            [
                ("lat", "lon"),
                ("latitude", "longitude"),
                ("pickuplat", "pickuplon"),
                ("pickuplat", "pickuplong"),
                ("deliverylat", "deliverylong"),
                ("droplat", "droplon"),
            ]
        )

        for lat_key, lon_key in candidates:
            if lat_key in obj and lon_key in obj and obj[lat_key] and obj[lon_key]:
                try:
                    return float(obj[lat_key]), float(obj[lon_key])
                except:
                    pass

        # Special case: nested 'pickup_location'
        if "pickup_location" in obj:
            return self.get_lat_lon(obj["pickup_location"])

        return 0.0, 0.0

    def get_order_kitchen(self, order: Dict[str, Any]) -> str:
        possible_keys = [
            "pickupcustomer",   # confirmed primary field in production orders
            "locationname",     # confirmed backup field in production orders
            "storename", "store_name",
            "restaurantname", "restaurant_name",
            "kitchenname", "kitchen_name",
            "partnername", "partner_name",
            "tenantname",
        ]
        for key in possible_keys:
            if key in order and order[key]:
                return str(order[key]).strip()
        return "Unknown"

    def _select_coverage_riders(
        self,
        valid_riders: List[Dict],
        clusters: List[Dict],
        target_count: int,
        rider_states: Dict,
    ) -> Set[int]:
        """
        Greedy set-cover: select the minimal set of riders that geographically
        covers all kitchen cluster centroids within MAX_PICKUP_DISTANCE_KM.

        Step 1: greedily pick rider that covers the most uncovered clusters.
        Step 2: pad up to target_count with remaining riders ordered by total
                distance to all cluster centroids (closest = most useful).

        Returns a set of rider IDs that form the preferred/pre-selected pool.
        """
        if not valid_riders or not clusters:
            return {r["id"] for r in valid_riders[:target_count]}

        # Sort by ID so greedy selection is deterministic regardless of
        # the order riders arrive in the API response.
        valid_riders = sorted(valid_riders, key=lambda r: r["id"])

        selected_ids: Set[int] = set()
        uncovered = list(range(len(clusters)))

        # Step 1 — greedy max-coverage selection
        while len(selected_ids) < target_count and uncovered:
            best_rider = None
            best_coverage = -1

            for r in valid_riders:
                if r["id"] in selected_ids:
                    continue
                coverage = sum(
                    1
                    for ci in uncovered
                    if self.haversine(
                        r["lat"], r["lon"],
                        clusters[ci]["centroid"][0], clusters[ci]["centroid"][1],
                    ) <= self.MAX_PICKUP_DISTANCE_KM
                )
                if coverage > best_coverage:
                    best_coverage = coverage
                    best_rider = r

            if not best_rider or best_coverage == 0:
                break

            selected_ids.add(best_rider["id"])
            # Remove clusters now covered by this rider
            uncovered = [
                ci for ci in uncovered
                if self.haversine(
                    best_rider["lat"], best_rider["lon"],
                    clusters[ci]["centroid"][0], clusters[ci]["centroid"][1],
                ) > self.MAX_PICKUP_DISTANCE_KM
            ]

        # Step 2 — pad with closest-aggregate riders up to target_count
        if len(selected_ids) < target_count:
            def total_dist_to_clusters(r):
                return sum(
                    self.haversine(
                        r["lat"], r["lon"],
                        c["centroid"][0], c["centroid"][1],
                    )
                    for c in clusters
                )

            remaining = sorted(
                [r for r in valid_riders if r["id"] not in selected_ids],
                key=total_dist_to_clusters,
            )
            for r in remaining:
                selected_ids.add(r["id"])
                if len(selected_ids) >= target_count:
                    break

        return selected_ids

    async def assign_orders(
        self,
        orders: List[Dict[str, Any]],
        riders: List[Dict[str, Any]],
        reshuffle: bool = False,
        fuel_charge: float = 2.5,
        base_pay: float = 0.0,
    ) -> tuple[Dict[int, List[Dict[str, Any]]], List[Dict[str, Any]]]:
        """
        ENHANCED: Cluster-Based Load-Balanced Assignment with Minimal Rider Selection.

        Strategy:
        1. Cluster orders by kitchen proximity
        2. Pre-select the MINIMUM set of riders needed to handle all orders
           (ceil(total_orders / max_orders_per_rider)) via greedy set-cover.
        3. FILL-BEFORE-SPILL: prefer already-started riders until ideal_load
           before activating a fresh rider — reduces rider count for small batches.
        4. Assign clusters to best-fit riders (proximity + workload balance)
        5. Rebalance if needed

        If reshuffle=True, controlled randomness is injected into rider scoring
        so that retrying the same input can explore alternative assignments.
        """
        from app.services.rider.rider_history_service import RiderHistoryService
        from app.services.rider.rider_state_manager import RiderStateManager
        from app.services.routing.clustering_service import ClusteringService

        # -- Load ML-tuned hyperparameters (or defaults on first run) ------
        self._load_config()
        # Use caller-supplied pricing so dynamic API rates flow into scoring
        self._fuel_rate = fuel_charge
        self._base_rider_cost = base_pay

        # 0. Prep
        assignments: Dict[int, List[Dict[str, Any]]] = defaultdict(list)
        unassigned_orders: List[Dict[str, Any]] = []
        rider_states = {}  # Track live load

        # 0a. KALMAN FILTER - Smooth rider GPS locations before scoring
        riders = smooth_rider_locations(list(riders))

        # 0b. KALMAN FILTER - Smooth order delivery coordinates before clustering
        orders = smooth_order_coordinates(list(orders))

        # 1. Parse and Filter Riders
        valid_riders = []
        BLOCKED_RIDERS = _BLOCKED_RIDERS_SET  # frozenset — single source of truth, O(1) lookup

        # Load Existing State (Persistence)
        state_mgr = RiderStateManager()

        for r in riders:
            # Robust ID Extraction
            rid_raw = r.get("userid") or r.get("riderid") or r.get("id") or r.get("_id")
            try:
                rid = int(rid_raw)
            except (ValueError, TypeError):
                continue

            if rid in BLOCKED_RIDERS:
                continue

            # Robust Status Check
            # Keep if: onduty (1, "1", True) OR status is active/idle/online
            is_onduty = str(r.get("onduty")) in ["1", "True"] or r.get("onduty") is True
            is_active = r.get("status") in ["active", "idle", "online"]

            if not (is_onduty or is_active):
                continue

            # Location
            lat, lon = self.get_lat_lon(r)

            # Fetch previous state to know if they are already busy
            p_state = state_mgr.get_rider_state(rid)

            # If rider has valid GPS, use it. If not, fallback to Last Drop or Home.
            if lat == 0 or lon == 0:
                if p_state["last_drop_lat"]:
                    lat, lon = p_state["last_drop_lat"], p_state["last_drop_lon"]
                else:
                    # Home Location Fallback (learned-augmented)
                    lat, lon = self.home_locations.get(rid, (0.0, 0.0))

            valid_riders.append({"id": rid, "lat": lat, "lon": lon, "obj": r})

            # Initialize rider state with existing workload
            existing_load = (
                p_state.get("minutes_remaining", 0) / 15
            )  # Convert minutes to order estimate

            rider_states[rid] = {
                "lat": lat,
                "lon": lon,
                "kitchens": set(),
                "count": int(existing_load),  # Start with existing workload
                "workload_score": existing_load,  # For prioritization
            }

        if not valid_riders:
            logger.warning(
                "No riders passed on-duty filter. Retrying with all available riders as emergency rescue..."
            )
            # If no on-duty riders, we take ANY rider provided by the API to ensure assignment
            for r in riders:
                rid = int(r.get("userid", 0))
                if rid in _BLOCKED_RIDERS_SET:
                    continue

                lat, lon = self.get_lat_lon(r)
                if lat == 0 or lon == 0:
                    lat, lon = self.home_locations.get(rid, (0.0, 0.0))

                if lat != 0:
                    valid_riders.append({"id": rid, "lat": lat, "lon": lon, "obj": r})
                    rider_states[rid] = {
                        "lat": lat,
                        "lon": lon,
                        "kitchens": set(),
                        "count": 0,
                        "workload_score": 0,
                    }

        if not valid_riders:
            logger.error("DANGER: Absolutely no riders available for assignment.")
            # Mark all as unassigned
            for o in orders:
                o["unassigned_reason"] = (
                    "No riders found (check partner online status)."
                )
                unassigned_orders.append(o)
            return assignments, unassigned_orders

        logger.info(f"Found {len(valid_riders)} active riders")

        # 2. CLUSTER ORDERS BY KITCHEN PROXIMITY
        clustering_service = ClusteringService()
        clusters = clustering_service.cluster_orders_by_kitchen(
            orders, max_cluster_radius_km=self.MAX_KITCHEN_DISTANCE_KM
        )

        # 2b. Tag each cluster with the set of kitchen names it covers (used
        # below for hard kitchen-ownership matching).
        for cluster in clusters:
            cluster["kitchen_names"] = set()
            for order in cluster["orders"]:
                cluster["kitchen_names"].add(self.get_order_kitchen(order))

        logger.info(f"Created {len(clusters)} order clusters")

        # 2c. MINIMAL RIDER PRE-SELECTION
        # Calculate the theoretical minimum number of riders needed so we don't
        # spray 40 orders across 9 riders when 4-5 would suffice.
        min_riders_needed = max(1, ceil(len(orders) / self.MAX_ORDERS_PER_RIDER))
        # Cap at available riders
        min_riders_needed = min(min_riders_needed, len(valid_riders))

        # Fair-share target: maximum orders any single rider should ideally carry.
        # Used as a per-batch cap so no rider consumes an entire large cluster
        # while other riders sit idle (the 15-16 vs 3-4 imbalance pattern).
        target_orders_per_rider = ceil(len(orders) / min_riders_needed)

        preferred_rider_ids = self._select_coverage_riders(
            valid_riders, clusters, min_riders_needed, rider_states
        )
        logger.info(
            f"[MinRider] {len(orders)} orders → min {min_riders_needed} riders needed "
            f"(capacity {self.MAX_ORDERS_PER_RIDER}/rider). "
            f"Pre-selected coverage pool: {preferred_rider_ids}"
        )

        # 3. ASSIGN CLUSTERS TO RIDERS (Fill-Before-Spill, Load-Balanced)
        for cluster_idx, cluster in enumerate(clusters):
            centroid_lat, centroid_lon = cluster["centroid"]
            cluster_orders = cluster["orders"]
            cluster_size = len(cluster_orders)

            logger.info(
                f"Assigning cluster {cluster_idx + 1}/{len(clusters)}: {cluster_size} orders at ({centroid_lat:.4f}, {centroid_lon:.4f})"
            )

            # Find best riders for this cluster
            candidate_riders = []

            # Calculate cluster-level profit features (Data Encoding: Engineered Features)
            cluster_profit = 0.0
            cluster_profit_density = 0.0
            cluster_geohash = _zone_key(centroid_lat, centroid_lon)

            for order in cluster_orders:
                o_lat = float(order.get("pickuplat", 0))
                o_lon = float(order.get("pickuplon", 0))
                dist = (
                    self.haversine(centroid_lat, centroid_lon, o_lat, o_lon)
                    if o_lat
                    else 2.0
                )

                profit_features = self.calculate_order_profit_features(order, dist)
                cluster_profit += profit_features["profit"]
                cluster_profit_density += profit_features["profit_density"]

            avg_profit_density = (
                cluster_profit_density / cluster_size if cluster_size > 0 else 0
            )

            # ── HARD KITCHEN OWNERSHIP (rider_preferences.py = source of truth) ──
            # Step 1: find which rider IDs are configured for this cluster's kitchens
            cluster_preferred_rids: set = set()
            for k_name in cluster["kitchen_names"]:
                k_lower = k_name.lower()
                for cfg_rid, cfg_prefs in self.rider_preferences.items():
                    for p in cfg_prefs:
                        if p.lower() in k_lower or k_lower in p.lower():
                            cluster_preferred_rids.add(cfg_rid)

            # Step 2: among those, which are actually in the active fleet with capacity?
            active_preferred_with_cap = {
                r["id"] for r in valid_riders
                if r["id"] in cluster_preferred_rids
                and rider_states[r["id"]]["count"] < self.MAX_ORDERS_PER_RIDER
            }

            # RULE: if at least one configured rider is available with capacity,
            # only that rider (or those riders) can be assigned this cluster's orders.
            # Non-configured riders are completely skipped.
            # If NO configured rider is available (all full / off-duty), open to all.
            kitchen_has_owner = bool(active_preferred_with_cap)

            if cluster_preferred_rids:
                logger.info(
                    f"  [Kitchen owner] cluster kitchens={list(cluster['kitchen_names'])} "
                    f"→ configured riders={cluster_preferred_rids} "
                    f"| active+capacity={active_preferred_with_cap} "
                    f"| hard_lock={'YES' if kitchen_has_owner else 'NO (overflow)'}"
                )

            for r in valid_riders:
                rid = r["id"]
                r_state = rider_states[rid]

                # ── HARD KITCHEN LOCK ────────────────────────────────────────
                # If this kitchen has an owner and this rider isn't the owner → skip
                if kitchen_has_owner and rid not in active_preferred_with_cap:
                    continue

                # Calculate distance to cluster centroid
                dist = self.haversine(
                    r_state["lat"], r_state["lon"], centroid_lat, centroid_lon
                )

                # Preference flag (SOFT: distance-limit bypass + preference bonus).
                # Uses learned-augmented prefs so riders who actually serve a
                # kitchen are steered there even if not in the curated config.
                prefs = self.soft_preferences.get(rid, [])
                has_preference = False
                for k_name in cluster["kitchen_names"]:
                    if any(
                        p.lower() in k_name.lower() or k_name.lower() in p.lower()
                        for p in prefs
                    ):
                        has_preference = True
                        break

                # Dynamic Limit: 6km default, 10km for preferred kitchens
                allowed_dist = self.MAX_PICKUP_DISTANCE_KM
                if has_preference:
                    allowed_dist = max(allowed_dist, 10.0)

                # Skip if too far
                if dist > allowed_dist:
                    continue

                # Calculate workload utilization (0.0 to 1.0)
                utilization = r_state["count"] / self.MAX_ORDERS_PER_RIDER

                # ── FILL-BEFORE-SPILL workload scoring ──────────────────────
                # Goal: fill a started rider to IDEAL_LOAD before routing orders
                # to a fresh rider. This naturally minimises active rider count.
                #
                #  • count == 0, not in pre-selected pool → NEW_RIDER_PENALTY
                #    (heavy cost to discourage opening a new rider unnecessarily)
                #  • count == 0, in pre-selected pool → neutral (no penalty)
                #  • 0 < count < IDEAL_LOAD → FILL_BONUS (negative = better score)
                #    bonus grows as the rider is closer to full
                #  • count >= IDEAL_LOAD → linear workload penalty (ML-tuned)
                # ─────────────────────────────────────────────────────────────
                # Use the dynamic NEW_RIDER_PENALTY computed in _load_config
                # (= WORKLOAD_PENALTY_WEIGHT × 1.5, always above any possible
                # workload score so non-preferred riders never open early).
                NEW_RIDER_PENALTY = self.NEW_RIDER_PENALTY
                FILL_BONUS_PER_SLOT = 8.0  # reward per empty slot below ideal_load

                is_preferred = rid in preferred_rider_ids
                current_count = r_state["count"]

                if current_count == 0:
                    if is_preferred:
                        workload_penalty = 0.0  # Pre-selected, neutral cost
                    else:
                        workload_penalty = NEW_RIDER_PENALTY  # Discourage extra riders
                elif current_count < self.IDEAL_LOAD:
                    # Strong incentive to fill an already-started rider
                    slots_remaining = self.IDEAL_LOAD - current_count
                    workload_penalty = -(slots_remaining * FILL_BONUS_PER_SLOT)
                else:
                    # Above ideal load — apply ML-tuned linear penalty
                    workload_penalty = utilization * self.WORKLOAD_PENALTY_WEIGHT

                distance_penalty = dist * self.DISTANCE_PENALTY_WEIGHT

                # Preference bonus (ML-tuned)
                preference_bonus = self.PREFERENCE_BONUS if has_preference else 0

                # Home zone bonus (ML-tuned). Uses learned-augmented home coords.
                h_lat, h_lon = self.home_locations.get(rid, (0.0, 0.0))
                home_bonus = 0
                if h_lat != 0:
                    home_dist = self.haversine(h_lat, h_lon, centroid_lat, centroid_lon)
                    if home_dist <= 4.0:
                        home_bonus = self.HOME_ZONE_BONUS_4KM
                    if home_dist <= 2.0:
                        home_bonus = self.HOME_ZONE_BONUS_2KM

                # PROFIT-AWARE SCORING (Data Encoding: Composite Cost Function)
                # Higher profit density clusters should get priority (lower score = better)
                profit_bonus = avg_profit_density * 10  # Scale up profit density signal

                # Geohash zone proximity bonus
                # The _simple_geohash produces a binary bit-string (only '0'/'1' chars).
                # precision=25 → 13 lon bits + 12 lat bits (alternating).
                # Global lon range (360°) / 2^13 = 0.044° ≈ 4.9 km per cell.
                # Global lat range (180°) / 2^12 = 0.044° ≈ 4.9 km per cell.
                # Full 25-char prefix comparison = exact same ~5km cell.
                h_geohash = _zone_key(h_lat, h_lon) if h_lat != 0 else ""
                geohash_match_bonus = 0
                if h_geohash and cluster_geohash not in ("", "unknown"):
                    if h_geohash == cluster_geohash:
                        geohash_match_bonus = -15  # Negative = better score (≈6km zone match)

                score = (
                    workload_penalty
                    + distance_penalty
                    + preference_bonus
                    + home_bonus
                    - profit_bonus
                    + geohash_match_bonus  # Profit-aware scoring
                )

                # RESHUFFLE: Add controlled noise so retries explore different riders
                if reshuffle:
                    noise = random.uniform(-15.0, 15.0)
                    score += noise

                candidate_riders.append(
                    {
                        "id": rid,
                        "score": score,
                        "distance": dist,
                        "utilization": utilization,
                        "current_load": r_state["count"],
                        # Store per-candidate bonuses so the re-scorer inside the
                        # inner while-loop can reconstruct the full score correctly
                        # (Bug fix: old re-scorer dropped these, reverting to
                        # plain distance+workload after the first batch assignment).
                        "_preference_bonus": preference_bonus,
                        "_home_bonus": home_bonus,
                        "_profit_bonus": profit_bonus,
                        "_geohash_match_bonus": geohash_match_bonus,
                    }
                )

            if not candidate_riders:
                logger.warning(f"No riders available for cluster {cluster_idx + 1}")
                for o in cluster_orders:
                    o["unassigned_reason"] = (
                        f"No riders within {self.MAX_PICKUP_DISTANCE_KM}km radius of kitchen."
                    )
                    unassigned_orders.append(o)
                continue

            # Sort by score (best first)
            candidate_riders.sort(key=lambda x: x["score"])

            # SMART DISTRIBUTION: Split cluster if needed
            remaining_orders = cluster_orders[:]

            while remaining_orders and candidate_riders:
                best_rider = candidate_riders[0]
                rid = best_rider["id"]
                r_state = rider_states[rid]

                # How many orders can this rider take?
                available_capacity = self.MAX_ORDERS_PER_RIDER - r_state["count"]

                if available_capacity <= 0:
                    # Rider is full, remove from candidates
                    candidate_riders.pop(0)
                    continue

                # Decide batch size
                # ── FAIR-SHARE CAP ────────────────────────────────────────────
                # Never give a rider more than (target_orders_per_rider - current)
                # in a single iteration. This prevents one rider consuming an
                # entire large cluster (e.g. 15 orders) while 3 others stay idle.
                # max(1, ...) ensures the loop always makes forward progress even
                # if the rider has slightly exceeded their target.
                _fairshare_remaining = max(1, target_orders_per_rider - r_state["count"])

                if best_rider["utilization"] < self.WORKLOAD_BALANCE_THRESHOLD:
                    # Rider has capacity — cap at fair share for balanced distribution
                    batch_size = min(available_capacity, len(remaining_orders), _fairshare_remaining)
                else:
                    # Rider is getting busy, be conservative (IDEAL_LOAD from ML)
                    batch_size = min(
                        self.IDEAL_LOAD - r_state["count"],
                        len(remaining_orders),
                        available_capacity,
                        _fairshare_remaining,
                    )
                    batch_size = max(1, batch_size)  # At least 1 order

                # Assign batch
                batch = remaining_orders[:batch_size]
                remaining_orders = remaining_orders[batch_size:]

                assignments[rid].extend(batch)

                # Update rider state
                r_state["count"] += len(batch)
                r_state["lat"] = centroid_lat
                r_state["lon"] = centroid_lon
                r_state["kitchens"].update(cluster["kitchen_names"])
                r_state["workload_score"] = r_state["count"] / self.MAX_ORDERS_PER_RIDER

                logger.info(
                    f"  -> Assigned {len(batch)} orders to Rider {rid} (load: {r_state['count']}/{self.MAX_ORDERS_PER_RIDER})"
                )

                # Re-sort candidates by updated scores (use ML-tuned weights, not hardcoded).
                # Only the rider that just received orders needs a score update; all others
                # keep their current scores (their workload didn't change).
                for candidate in candidate_riders:
                    if candidate["id"] == rid:
                        new_count = rider_states[candidate["id"]]["count"]
                        new_util = new_count / self.MAX_ORDERS_PER_RIDER
                        candidate["utilization"] = new_util
                        candidate["current_load"] = new_count

                        # Reapply fill-before-spill workload component
                        if new_count == 0:
                            _wp = 0.0 if candidate["id"] in preferred_rider_ids else self.NEW_RIDER_PENALTY
                        elif new_count < self.IDEAL_LOAD:
                            _wp = -((self.IDEAL_LOAD - new_count) * 8.0)
                        else:
                            _wp = new_util * self.WORKLOAD_PENALTY_WEIGHT

                        # Bug fix: restore ALL scoring components, not just workload+distance.
                        # The previous code dropped preference_bonus, home_bonus, profit_bonus,
                        # and geohash_match_bonus, reverting to plain distance+workload
                        # for every cluster iteration after the first batch was assigned.
                        candidate["score"] = (
                            _wp
                            + candidate["distance"] * self.DISTANCE_PENALTY_WEIGHT
                            + candidate["_preference_bonus"]
                            + candidate["_home_bonus"]
                            - candidate["_profit_bonus"]
                            + candidate["_geohash_match_bonus"]
                        )

                candidate_riders.sort(key=lambda x: x["score"])

            # If any orders left in the cluster after exhaustion of candidates
            if remaining_orders:
                # Instead of giving up, keep them in a pool for mandatory assignment
                unassigned_orders.extend(remaining_orders)

        # 4. EMERGENCY MANDATORY ASSIGNMENT (Ensures 0 unassigned if riders exist)
        if unassigned_orders and valid_riders:
            logger.info(
                f"[ALERT] Starting Emergency Mandatory Assignment for {len(unassigned_orders)} orders..."
            )
            force_pool = unassigned_orders[:]
            unassigned_orders.clear()

            for o in force_pool:
                # Determine pickup location
                o_lat, o_lon = self.get_lat_lon(o, prefix="pickup")
                if o_lat == 0:
                    o["unassigned_reason"] = "Could not geolocate order (0,0)."
                    unassigned_orders.append(o)
                    continue

                # Find the 'least bad' rider (Closest + Balanced Load)
                best_emergency_rider = None
                best_emergency_score = float("inf")

                for r in valid_riders:
                    rid = r["id"]
                    r_state = rider_states[rid]

                    dist = self.haversine(r_state["lat"], r_state["lon"], o_lat, o_lon)
                    # For emergency: Distance is important, but load prevents one rider taking EVERYTHING
                    # Score = distance + ML-tuned penalty per existing order
                    e_score = dist + (r_state["count"] * self.EMERGENCY_LOAD_PENALTY)

                    if e_score < best_emergency_score:
                        best_emergency_score = e_score
                        best_emergency_rider = rid

                if best_emergency_rider:
                    assignments[best_emergency_rider].append(o)
                    rider_states[best_emergency_rider]["count"] += 1
                    logger.info(
                        f"   Force-Assigned order {o.get('orderid')} to Rider {best_emergency_rider} (Score: {best_emergency_score:.2f})"
                    )
                else:
                    unassigned_orders.append(o)

        # 5. FINAL REBALANCING (Optional)
        # Check if any rider is overloaded while others are idle
        self._rebalance_workload(assignments, rider_states, valid_riders)

        # 6. Commit State and History
        self._post_process(assignments, rider_states, state_mgr)

        # ML event logging happens once, in the /riderassign endpoint after
        # Phase-0 history merge + solo consolidation — not here — so every
        # request produces exactly one assignment_ml_log row (see
        # optimization.py::_bg_log_assignment).

        # Log final distribution (use r_orders to avoid shadowing the outer `orders` list)
        logger.info("=" * 50)
        logger.info("FINAL ASSIGNMENT DISTRIBUTION:")
        for rid, r_orders in sorted(assignments.items()):
            logger.info(f"  Rider {rid}: {len(r_orders)} orders")

        if unassigned_orders:
            logger.warning(
                f"  [ALERT] STILL UNASSIGNED: {len(unassigned_orders)} (Reason: No riders online or invalid coords)"
            )
        else:
            logger.info("  [OK] ALL ORDERS ASSIGNED SUCCESSFULLY")
        logger.info("=" * 50)

        return assignments, unassigned_orders

    def _rebalance_workload(
        self, assignments: Dict[int, List], rider_states: Dict, valid_riders: List
    ):
        """
        Rebalance if workload is heavily skewed.
        Move orders from overloaded riders to idle ones if possible.
        """
        if not assignments:
            return

        # Calculate average load — divide by ACTIVE riders only (those with ≥1 order),
        # not all valid riders. Using all riders artificially lowers avg_load and causes
        # false-positive over/under-load detection.
        total_orders = sum(len(rl) for rl in assignments.values())
        active_rider_count = sum(1 for r in valid_riders if rider_states[r["id"]]["count"] > 0)
        avg_load = total_orders / active_rider_count if active_rider_count > 0 else 0

        # Find overloaded and underutilized riders
        overloaded = []
        underutilized = []

        for r in valid_riders:
            rid = r["id"]
            load = rider_states[rid]["count"]

            if load > avg_load * 1.5 and load > self.IDEAL_LOAD:  # 50% above average
                overloaded.append(rid)
            elif load < avg_load * 0.5:  # 50% below average
                underutilized.append(rid)

        if not overloaded or not underutilized:
            return

        logger.info(
            f"Rebalancing: {len(overloaded)} overloaded, {len(underutilized)} underutilized riders"
        )

        # Try to move orders from overloaded to underutilized
        for over_rid in overloaded:
            over_orders = assignments[over_rid]
            over_state = rider_states[over_rid]

            # Try to offload some orders
            for under_rid in underutilized:
                under_state = rider_states[under_rid]
                under_capacity = self.MAX_ORDERS_PER_RIDER - under_state["count"]

                if under_capacity <= 0:
                    continue

                # Find orders that are closer to underutilized rider
                transferable = []
                for order in over_orders:
                    o_lat, o_lon = self.get_lat_lon(order, prefix="pickup")
                    if o_lat == 0:
                        continue

                    dist_to_under = self.haversine(
                        under_state["lat"], under_state["lon"], o_lat, o_lon
                    )
                    dist_to_over = self.haversine(
                        over_state["lat"], over_state["lon"], o_lat, o_lon
                    )

                    # Transfer if underutilized rider is closer or similar distance
                    if (
                        dist_to_under <= self.MAX_PICKUP_DISTANCE_KM
                        and dist_to_under <= dist_to_over * 1.2
                    ):
                        transferable.append(order)

                if transferable:
                    # Transfer up to capacity.
                    # over_state["count"] - IDEAL_LOAD is the surplus above ideal;
                    # guard with max(1, ...) so we always move at least one order
                    # and never get a zero/negative transfer_count.
                    surplus = max(1, over_state["count"] - self.IDEAL_LOAD)
                    transfer_count = min(
                        len(transferable),
                        under_capacity,
                        surplus,
                    )
                    transfer_batch = transferable[:transfer_count]

                    # Move orders
                    for order in transfer_batch:
                        over_orders.remove(order)
                        assignments[under_rid].append(order)

                    # Update states
                    over_state["count"] -= len(transfer_batch)
                    under_state["count"] += len(transfer_batch)

                    logger.info(
                        f"  Rebalanced: {len(transfer_batch)} orders from Rider {over_rid} -> {under_rid}"
                    )

    def _post_process(self, assignments, rider_states, state_mgr):
        """Update History and Persistence."""
        from app.services.rider.rider_history_service import RiderHistoryService

        history_service = RiderHistoryService()
        ts = time.time()

        for rid, rider_orders in assignments.items():
            if not rider_orders:
                continue

            # Calculate actual cumulative distance for this rider instead of
            # the old hardcoded 5.0 km placeholder.
            total_km = 0.0
            for o in rider_orders:
                try:
                    total_km += float(o.get("actualkms") or o.get("kms") or 0)
                except (TypeError, ValueError):
                    pass

            history_service.update_rider_stats(rid, total_km, len(rider_orders))

            st = rider_states[rid]
            # ETA per order ≈ 15 min (used for workload decay in RiderStateManager)
            state_mgr.states[rid] = {
                "minutes_remaining": len(rider_orders) * 15,
                "last_drop_lat": st["lat"],
                "last_drop_lon": st["lon"],
                "active_kitchens": st["kitchens"],
                "last_updated_ts": ts,
            }

        state_mgr._save_states()
