"""Greedy group-preserving batching.

Packs items into batches of at most ``max_batch_size`` without ever splitting
a group (items sharing a key) across batches; a single group larger than the
cap becomes its own oversized batch. Shared by the address postcode batcher
(``domain/addresses/postcode_batching.py``) and the Modelling Run distributor
(``backend/app/modelling/batching.py``).
"""

from collections.abc import Callable, Hashable, Iterable, Iterator
from typing import TypeVar

T = TypeVar("T")


def iter_grouped_batches(
    items: Iterable[T],
    *,
    key: Callable[[T], Hashable],
    max_batch_size: int,
) -> Iterator[list[T]]:
    if max_batch_size < 1:
        raise ValueError("max_batch_size must be >= 1")

    groups: dict[Hashable, list[T]] = {}
    for item in items:
        groups.setdefault(key(item), []).append(item)

    buffer: list[T] = []
    for group in groups.values():
        # Oversize single-key group: flush the buffer first, then dispatch
        # the group as its own batch.
        if len(group) >= max_batch_size:
            if buffer:
                yield buffer
                buffer = []
            yield group
            continue

        # Adding this group would overflow: flush the buffer before appending.
        if len(buffer) + len(group) > max_batch_size:
            yield buffer
            buffer = []

        buffer.extend(group)

    if buffer:
        yield buffer