"""Format Response Data."""

from collections import defaultdict


def prepare_dataload_response(entity_ids, entity_list, key_field_name):
    """Prepare dataload response.

    Prepare dataload response so all the ids will be handled:
    [
        {
            'data': {
                'abacus_event_id': 123,
                'event_name': 321,
                ...
            }
        },
        {
            'data': None
        }
    ]
    """
    entity_by_id = {entity.get(key_field_name): entity for entity in entity_list}

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

    return result


def prepare_dataload_with_data_as_list_response(
    entity_ids, entity_list, key_field_name
):
    """Prepare a one-to-many dataload response.

    Like prepare_dataload_response, but each id maps to the LIST of records that
    carry it (or None when there are none), so a single id can return many rows:
    [
        {'data': [{'abacus_event_id': 1, ...}, {'abacus_event_id': 2, ...}]},
        {'data': None},
    ]

    Groups by key in a single pass so shaping is O(ids + records); filtering the
    whole list once per id would be O(ids x records) and a cliff for a
    high-fan-out endpoint.
    """
    entities_by_id = defaultdict(list)
    for entity in entity_list:
        entities_by_id[entity.get(key_field_name)].append(entity)

    return [{'data': entities_by_id.get(entity_id) or None} for entity_id in entity_ids]