# app/scheduling/distributor.py

import math
from app.assignment.rules import (
    MAX_PRODUCT_QUANTITY,
    MAX_DIFFICULT_QUANTITY,
    MAX_FRIDGE_QUANTITY,
    MAKE_ALONE_PRODUCTS,
    MAX_DIFFICULT_PRODUCTS,
)
from app.assignment.assigner import OrderAssigner


class OrderDistributor:
    """
    Heuristic load-based distributor inspired by Least Load First scheduling.
    """

    def __init__(self, product_costs: dict, waiters: list[str]):
        self.product_costs = product_costs
        self.groups = list(product_costs.keys())

        self.product_cost = {
            product: cost
            for group in product_costs
            for product, cost in product_costs[group].items()
        }

        self.waiters = waiters
        self.waiter_load: dict[str, float] = {w: 0 for w in waiters}
        self.order_assigner = OrderAssigner()

    def distribute_single_order(self, algorithm, order: dict[str, int], last_order_waiters: list[str]) -> dict:
        """
        Returns:
            {
                product_name: {
                    waiter_name: quantity
                }
            }
        """

        # --- initialize state ---
        waiter_products = {w: {} for w in self.waiters}
        waiter_order_load = {w: 0 for w in self.waiters}
        estimated_load = self.waiter_load.copy()

        # --- assignment reusable function ---
        def assign(waiter: str, product: str, qty: int):
            if qty <= 0:
                return

            cost = self.product_cost.get(product, 0)
            load = cost + (cost / 2)* (qty - 1) # ∆Li(p, q) = c(p) + (c(p)/2) * (q − 1),

            estimated_load[waiter] += load # Li ← Li + ∆Li(p, q)
            waiter_order_load[waiter] += load
            waiter_products[waiter][product] = (
                waiter_products[waiter].get(product, 0) + qty
            )

        for group in self.groups:
            current_waiter = self.order_assigner.assign_order(
                algorithm=algorithm,
                waiter_products=waiter_products,
                waiter_loads=estimated_load,
                waiter_order_loads=waiter_order_load,
                last_order_waiters=last_order_waiters,
            )
            units_assigned = {w: 0 for w in self.waiters}

            for product, quantity in order.items():
                if product not in self.product_costs[group]:
                    continue

                limit = self._group_limit(group, product)

                if units_assigned[current_waiter] + quantity >= limit:
                    split = math.floor(quantity / 2)
                    assign(current_waiter, product, split)
                    units_assigned[current_waiter] += split

                    remaining = quantity - split
                    if remaining > 0:
                        next_waiter = self.order_assigner.assign_order(
                            algorithm=algorithm,
                            waiter_products=waiter_products,
                            waiter_loads=estimated_load,
                            waiter_order_loads=waiter_order_load,
                            last_order_waiters=last_order_waiters,
                            exclude=[current_waiter],
                        )
                        assign(next_waiter, product, remaining)
                        units_assigned[next_waiter] += remaining
                else:
                    assign(current_waiter, product, quantity)
                    units_assigned[current_waiter] += quantity

                if self._should_switch_waiter(group, product, waiter_products[current_waiter].get(product, 0)):
                    current_waiter = self.order_assigner.assign_order(
                        algorithm=algorithm,
                        waiter_products=waiter_products,
                        waiter_loads=estimated_load,
                        waiter_order_loads=waiter_order_load,
                        last_order_waiters=last_order_waiters,
                    )

        return self._build_distribution(order, waiter_products), estimated_load
    
    # ---------- internal helpers ----------

    def _group_limit(self, group: str, product: str) -> int:
        if group == "fridges":
            return MAX_FRIDGE_QUANTITY
        if product in MAKE_ALONE_PRODUCTS:
            return MAX_DIFFICULT_QUANTITY
        return MAX_PRODUCT_QUANTITY

    def _should_switch_waiter(self, group: str, product: str, qty: int) -> bool:
        return (
            group == "alone_make_products"
            or (product in MAX_DIFFICULT_PRODUCTS
            and qty >= 2)
        )

    @staticmethod
    def _build_distribution(order, waiter_products):
        distribution = {}
        for product in order:
            for waiter, products in waiter_products.items():
                if product in products:
                    distribution.setdefault(product, {})
                    distribution[product][waiter] = products[product]
        return distribution
