"""Format Response Data."""

from collections import defaultdict
from collections.abc import Callable, Hashable, Iterable
from typing import Any

Record = dict[str, Any]
DataloadEntry = dict[str, Record | None]
DataloadListEntry = dict[str, list[Record] | None]


def prepare_dataload_response(
    entity_ids: list[int], entity_list: list[Record], key_field_name: str
) -> list[DataloadEntry]:
    """Prepare dataload response.

    Prepare dataload response so all the ids will be handled:
    [
        {
            'data': {
                'contract_lifecycle_schedule_id': 123,
                'contract_id': 321,
                ...
            }
        },
        {
            'data': None
        }
    ]

    Args:
        entity_ids(list): a list of entity identifiers
        entity_list(list): a list of records pertaining to an entity
        key_field_name(str): an identifier field that can be used to determine
            whether a record exists
    """
    entity_by_id = {entity.get(key_field_name): entity for entity in entity_list}

    result: list[DataloadEntry] = [
        {'data': entity_by_id.get(entity_id)} for entity_id in entity_ids
    ]

    return result


def prepare_dataload_with_data_as_list_response(
    entity_ids: list[int], entity_list: list[Record], key_field_name: str
) -> list[DataloadListEntry]:
    """Prepare dataload response containing list.

    Prepare dataload response so all the ids will be handled:
    [
        {
            'data': [
                {
                    'account_id': 123,
                    'contract_id': 321,
                    ...
                },
                {
                    'account_id': 123,
                    'contract_id': 432,
                    ...
                },
            ]
        },
        {
            'data': None
        }
    ]

    Args:
        entity_ids(list): a list of entity identifiers
        entity_list(list): a nested list of records pertaining to an entity
        key_field_name(str): an identifier field that can be used to determine
            whether a record exists
    """
    # Group by key in a single pass so shaping is O(ids + records); the prior
    # per-id filter re-scanned the whole record list once per id (and again per
    # duplicate id), which is O(ids x records) and becomes a synchronous cliff
    # once a high-fan-out endpoint migrates onto this shared path.
    entities_by_id: defaultdict[Any, list[Record]] = defaultdict(list)
    for entity in entity_list:
        entities_by_id[entity.get(key_field_name)].append(entity)

    results: list[DataloadListEntry] = [
        {'data': entities_by_id.get(entity_id) or None} for entity_id in entity_ids
    ]

    return results


def prepare_dataload_grouped_response(
    entity_ids: list[int],
    pairs: Iterable[Any],
    *,
    group_key: Callable[[Any], Any],
    child: Callable[[Any], Any],
    serialize_group: Callable[[list[Any]], Any],
    dedup_key: Callable[[Any], Hashable] | None = None,
    sort_key: Callable[[Any], Any] | None = None,
) -> list[DataloadListEntry]:
    """One-to-many dataload shaper for children that don't carry the parent key.

    Unlike ``prepare_dataload_with_data_as_list_response`` (which groups
    already-dumped records by a key field present on each record), this groups
    link/junction rows by a parent id derived from each pair, so the child record
    need not reference its parent.

    For each requested id, in order: collect that id's children, optionally dedup
    by ``dedup_key`` (first occurrence wins) and sort by ``sort_key``, then
    ``serialize_group`` the list into the ``data`` payload (``None`` when empty).
    Single pass over the pairs, so O(pairs + ids).

    Args:
        entity_ids: requested parent ids, in response order.
        pairs: rows linking a parent to a child (e.g. junction rows).
        group_key: pair -> parent id.
        child: pair -> the child object to collect.
        serialize_group: children list -> the serialized ``data`` payload.
        dedup_key: optional child -> hashable identity for de-duplication.
        sort_key: optional child -> sort key.
    """
    children_by_id: defaultdict[Any, list[Any]] = defaultdict(list)
    seen_by_id: defaultdict[Any, set[Hashable]] = defaultdict(set)
    for pair in pairs:
        parent_id = group_key(pair)
        record = child(pair)
        if dedup_key is not None:
            identity = dedup_key(record)
            if identity in seen_by_id[parent_id]:
                continue
            seen_by_id[parent_id].add(identity)
        children_by_id[parent_id].append(record)

    def payload(parent_id: Any) -> Any:
        children = children_by_id.get(parent_id)
        if not children:
            return None
        if sort_key is not None:
            children = sorted(children, key=sort_key)
        return serialize_group(children)

    return [{'data': payload(entity_id)} for entity_id in entity_ids]