# app/scheduling/load.py

from __future__ import annotations

import random
from dataclasses import dataclass
from typing import Optional


@dataclass
class AssignmentInput:
    """Lightweight container for the state needed by selection algorithms."""
    waiter_products: dict[str, dict]
    waiter_loads: dict[str, float]
    waiter_order_loads: dict[str, float]
    last_order_waiters: Optional[list[str]] = None
    exclude: Optional[list[str]] = None


class OrderAssigner:

    def __init__(self, seed: int = 42) -> None:
        self._rr_idx: int = 0  # Round-robin pointer (global over all waiters)
        self._rng = random.Random(seed)  # Deterministic RNG for reproducible experiments
        
    def assign_order(
        self,
        algorithm: str,
        waiter_products: dict[str, dict],
        waiter_loads: dict[str, float],
        waiter_order_loads: dict[str, float],
        last_order_waiters: list[str] | None = None,
        exclude: list[str] | None = None,
    ) -> str:
        """
        Selects and returns the chosen waiter name according to the requested algorithm.
        """
        inp = AssignmentInput(
            waiter_products=waiter_products,
            waiter_loads=waiter_loads,
            waiter_order_loads=waiter_order_loads,
            last_order_waiters=last_order_waiters,
            exclude=exclude,
        )

        match algorithm:
            case "round-robin":
                return self.round_robin(inp)
            case "least-loaded":
                return self.least_loaded(inp)
            case "least-loaded-speed":
                return self.least_loaded_speed(inp)
            case "power-of-two-choices":
                return self.power_of_two_choices(inp)
            case "custom-1":
                return self.custom_1(inp)
            case "custom-2":
                return self.custom_2(inp)
            case _:
                raise ValueError(f"Unknown algorithm: {algorithm}")

    # -------------------------
    # Algorithms (return waiter name)
    # -------------------------

    def round_robin(self, inp: AssignmentInput) -> str:
        """
        Pure global round-robin selection:
        - Cycles through the full waiter set in a fixed order.
        - If the next waiter is excluded, advances until it finds an eligible one.
        """
        exclude_set = set(inp.exclude or [])

        # Global fixed order over ALL waiters
        all_waiters = sorted(inp.waiter_loads.keys())
        n = len(all_waiters)
        if n == 0:
            raise ValueError("No waiters available.")

        # Try up to n times to find a non-excluded waiter
        for _ in range(n):
            w = all_waiters[self._rr_idx % n]
            self._rr_idx = (self._rr_idx + 1) % n  # Always advance globally
            if w not in exclude_set:
                return w

        raise ValueError("No eligible waiters (all excluded).")

    def least_loaded(self, inp: AssignmentInput) -> str:
        """Pure least-loaded selection over current accumulated loads."""
        exclude_set = set(inp.exclude or [])
        candidates = [w for w in inp.waiter_loads.keys() if w not in exclude_set]
        if not candidates:
            raise ValueError("No eligible waiters (all excluded).")
        return min(candidates, key=lambda w: inp.waiter_loads[w])

    def least_loaded_speed(self, inp: AssignmentInput) -> str:
        """
        Least-loaded with 'speed' heuristics (same spirit as your previous function):
        - Prefer waiters NOT assigned in the last order (avoid repeating).
        - Strongly prefer waiters NOT already used in the current order (avoid duplicates).
        - Fallback to least-loaded among allowed waiters.
        """
        exclude_set = set(inp.exclude or [])
        last_set = set(inp.last_order_waiters or [])
        # Sort by global load ascending
        sorted_waiters = sorted(
            ((w, load) for w, load in inp.waiter_loads.items() if w not in exclude_set),
            key=lambda x: x[1],
        )
        if not sorted_waiters:
            raise ValueError("No eligible waiters (all excluded).")

        # 1) Best case: not in last order AND unused in current order
        for w, _ in sorted_waiters:
            if w in last_set:
                continue
            if inp.waiter_products.get(w, {}) == {}:
                return w

        # 2) Next: unused in current order (even if in last order)
        for w, _ in sorted_waiters:
            if inp.waiter_products.get(w, {}) == {}:
                return w

        # 3) Fallback: least-loaded overall
        return sorted_waiters[0][0]

    def power_of_two_choices(self, inp: AssignmentInput) -> str:
        """
        Power of two choices (deterministic with fixed seed):
        - Pick 2 random eligible waiters and choose the one with smaller load.
        """
        exclude_set = set(inp.exclude or [])
        candidates = [w for w in inp.waiter_loads.keys() if w not in exclude_set]

        if len(candidates) == 0:
            raise ValueError("No eligible waiters (all excluded).")
        if len(candidates) == 1:
            return candidates[0]

        a, b = self._rng.sample(candidates, 2)
        return a if inp.waiter_loads[a] <= inp.waiter_loads[b] else b

    def custom_1(self, inp: AssignmentInput) -> str:
        """
        Custom-1:
        Weighted score combining:
        - global accumulated load (fairness across the shift)
        - load already assigned in the current order (avoid concentrating a single order)
        - penalty for reusing a waiter already used in this order
        """
        exclude_set = set(inp.exclude or [])
        candidates = [w for w in inp.waiter_loads if w not in exclude_set]
        if not candidates:
            raise ValueError("No eligible waiters (all excluded).")

        # Normalization to prevent scale dominance
        max_global = max((inp.waiter_loads[w] for w in candidates), default=1.0) or 1.0
        max_order  = max((inp.waiter_order_loads[w] for w in candidates), default=1.0) or 1.0

        def score(w: str) -> float:
            global_norm = inp.waiter_loads[w] / max_global
            order_norm  = inp.waiter_order_loads[w] / max_order

            reused_penalty = 0.25 if inp.waiter_products.get(w, {}) else 0.0
            return 0.7 * global_norm + 0.3 * order_norm + reused_penalty

        return min(candidates, key=score)

    def custom_2(self, inp: AssignmentInput) -> str:
        """
        Custom-2 (Targeted Spread):
        Goal: maximize the number of distinct waiters engaged early in each order,
        while still keeping global fairness.

        Priority:
        1) Prefer waiters not yet used in the current order.
        2) Within those, prefer waiters not used in the last order.
        3) Use a blended score (global + order load) as tie-breaker / fallback.
        """
        
        exclude_set = set(inp.exclude or [])
        last_set = set(inp.last_order_waiters or [])


        eligible = [w for w in inp.waiter_loads.keys() if w not in exclude_set]
        if not eligible:
            raise ValueError("No eligible waiters (all excluded).")

        unused_in_order = [w for w in eligible if inp.waiter_products.get(w, {}) == {}]
        unused_and_not_last = [w for w in unused_in_order if w not in last_set]

        # Prepare normalization for fair scoring
        max_global = max((inp.waiter_loads[w] for w in eligible), default=1.0) or 1.0
        max_order  = max((inp.waiter_order_loads[w] for w in eligible), default=1.0) or 1.0

        def blended_score(w: str) -> float:
            # Lower is better
            global_norm = inp.waiter_loads[w] / max_global
            order_norm  = inp.waiter_order_loads[w] / max_order
            return 0.6 * global_norm + 0.4 * order_norm

        # 1) Spread: unused in current order, and avoid last order if possible
        if unused_and_not_last:
            return min(unused_and_not_last, key=blended_score)

        # 2) If can't avoid last order, still spread within the order
        if unused_in_order:
            return min(unused_in_order, key=blended_score)

        # 3) Everyone already used in this order: minimize additional concentration
        # Here we prioritize the waiter with the smallest order load first, then global.
        return min(eligible, key=lambda w: (inp.waiter_order_loads[w], inp.waiter_loads[w]))

