"""Overcooked order queue support.

Provides order queue tick logic and table builders. All pot/container
mechanics (interaction branches, tick handler, extra state, static tables)
are now auto-generated by the Container/Recipe declarative system in
``cogrid.core.containers`` and composed by ``cogrid.core.autowire``.

This module contains only order-queue-specific code:

- ``_build_order_tables`` -- build order queue static tables from config
- ``order_queue_tick`` -- advance order timers, handle expiry and spawning
- ``build_order_extra_state`` -- initialize order queue arrays
"""

from cogrid.backend import xp


def build_order_extra_state(order_config):
    """Initialize order queue arrays from config.

    Parameters
    ----------
    order_config : dict
        Must contain ``"max_active"``.

    Returns:
    -------
    dict
        Order queue arrays keyed with ``overcooked.`` prefix.
    """
    import numpy as _np

    max_active = order_config["max_active"]
    return {
        "overcooked.order_recipe": _np.full((max_active,), -1, dtype=_np.int32),
        "overcooked.order_timer": _np.zeros((max_active,), dtype=_np.int32),
        "overcooked.order_n_expired": _np.int32(0),
    }


def order_queue_tick(state, scope_config):
    """Advance order timers: countdown, expiry, and probabilistic spawning.

    Each step, for every empty order slot, independently samples from a
    categorical distribution over ``{no_order, recipe_0, recipe_1, ...}``
    using per-recipe spawn probabilities from ``order_spawn_probs``.

    Skips silently if no order arrays are present in extra_state.

    Parameters
    ----------
    state : EnvState
        Current state with optional order arrays in extra_state.
    scope_config : dict
        Must contain ``"static_tables"`` with ``order_spawn_probs``,
        ``order_time_limit``, and ``order_max_active``.

    Returns:
    -------
    EnvState
        Updated state (unchanged if no order arrays present).
    """
    import dataclasses

    from cogrid.backend._dispatch import get_backend

    if "overcooked.order_recipe" not in state.extra_state:
        return state

    static_tables = scope_config.get("static_tables", {})
    new_extra = dict(state.extra_state)

    order_recipe = new_extra["overcooked.order_recipe"]
    order_timer = new_extra["overcooked.order_timer"]
    order_n_expired = new_extra["overcooked.order_n_expired"]

    order_time_limit = static_tables["order_time_limit"]
    spawn_probs = static_tables["order_spawn_probs"]  # (n_recipes,)
    max_active = order_recipe.shape[0]

    # (1) Decrement active order timers
    active = order_recipe != -1
    new_order_timer = xp.where(active, order_timer - 1, order_timer)

    # (2) Clear expired orders (timer reached 0)
    expired = active & (new_order_timer <= 0)
    order_recipe = xp.where(expired, -1, order_recipe)
    new_order_timer = xp.where(expired, 0, new_order_timer)
    order_n_expired = order_n_expired + xp.sum(expired).astype(xp.int32)

    # (3) Probabilistic spawn per empty slot
    cdf = xp.cumsum(spawn_probs)  # (n_recipes,)

    if get_backend() == "jax":
        import jax

        key, subkey = jax.random.split(state.rng_key)
        rolls = jax.random.uniform(subkey, shape=(max_active,))
    else:
        import numpy as _np

        rolls = xp.array(_np.random.default_rng().uniform(size=(max_active,)), dtype=xp.float32)
        key = state.rng_key

    # hit[i, r] = True if slot i's roll falls in recipe r's CDF bucket
    hit = rolls[:, None] < cdf[None, :]  # (max_active, n_recipes)
    any_hit = xp.any(hit, axis=1)  # did this slot spawn?
    sampled_recipe = xp.argmax(hit.astype(xp.int32), axis=1)

    empty = order_recipe == -1
    spawn_mask = empty & any_hit

    order_recipe = xp.where(spawn_mask, sampled_recipe, order_recipe)
    new_order_timer = xp.where(spawn_mask, order_time_limit, new_order_timer)

    new_extra["overcooked.order_recipe"] = order_recipe
    new_extra["overcooked.order_timer"] = new_order_timer
    new_extra["overcooked.order_n_expired"] = order_n_expired

    return dataclasses.replace(state, extra_state=new_extra, rng_key=key)


def _build_order_tables(order_config, recipe_results):
    """Build order queue static tables from config dict.

    Parameters
    ----------
    order_config : dict or None
        If None, orders are disabled. Otherwise must contain
        ``spawn_probs`` (dict mapping result name to probability),
        ``max_active`` (default 3), ``time_limit`` (default 200).
    recipe_results : list[str]
        Ordered list of recipe result names (e.g. ``["onion_soup", "tomato_soup"]``)
        matching the recipe index ordering from ``Pot.recipes``.

    Returns:
    -------
    dict
        Keys: ``order_enabled``, and when enabled also
        ``order_spawn_probs``, ``order_max_active``, ``order_time_limit``.
    """
    import numpy as _np

    if order_config is None:
        return {"order_enabled": False}

    max_active = order_config.get("max_active", 3)
    time_limit = order_config.get("time_limit", 200)
    spawn_probs_dict = order_config["spawn_probs"]

    n_recipes = len(recipe_results)
    spawn_probs = _np.zeros(n_recipes, dtype=_np.float32)
    for name, prob in spawn_probs_dict.items():
        idx = recipe_results.index(name)
        spawn_probs[idx] = prob

    assert _np.all(spawn_probs >= 0), "All spawn probabilities must be >= 0"
    assert float(_np.sum(spawn_probs)) <= 1.0 + 1e-6, (
        f"spawn_probs must sum to <= 1.0, got {float(_np.sum(spawn_probs))}"
    )

    return {
        "order_enabled": True,
        "order_spawn_probs": spawn_probs,
        "order_max_active": _np.int32(max_active),
        "order_time_limit": _np.int32(time_limit),
    }


# ---------------------------------------------------------------------------
# Order queue HUD rendering
# ---------------------------------------------------------------------------

# Recipe index → display color (RGB).
_RECIPE_COLORS = [
    (255, 200, 50),  # recipe 0 (onion soup) – golden yellow
    (220, 50, 50),  # recipe 1 (tomato soup) – red
]
_EMPTY_COLOR = (80, 80, 80)
_BG_COLOR = (50, 50, 50)
_URGENT_COLOR = (255, 80, 80)


def build_order_hud_fn(order_config):
    """Return a HUD function that extracts order bar specs from env state.

    The returned function has signature::

        hud_fn(env_state, scope_config) -> list[dict] | None

    Each dict has keys ``progress`` (0-1 float), ``color`` (RGB tuple),
    and ``bg_color`` (RGB tuple).
    """
    time_limit = order_config.get("time_limit", 200)

    def hud_fn(env_state, scope_config):
        import numpy as _np

        extra = env_state.extra_state
        order_recipe = extra.get("overcooked.order_recipe")
        if order_recipe is None:
            return None

        order_timer = extra["overcooked.order_timer"]
        bars = []
        for i in range(len(order_recipe)):
            recipe_idx = int(_np.array(order_recipe[i]))
            timer = int(_np.array(order_timer[i]))
            if recipe_idx < 0:
                bars.append({"progress": 0, "color": _EMPTY_COLOR, "bg_color": _BG_COLOR})
            else:
                progress = timer / time_limit
                color = _RECIPE_COLORS[recipe_idx % len(_RECIPE_COLORS)]
                if progress < 0.2:
                    color = _URGENT_COLOR
                bars.append({"progress": progress, "color": color, "bg_color": _BG_COLOR})
        return bars

    return hud_fn
