"""Format Response Data."""


def prepare_dataload_response(
    entity_ids: list,
    entity_list: list,
    key_field_name: str
):
    """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 = [
        {
            '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: list,
    entity_list: list,
    key_field_name: str
):
    """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
    """
    results = [
        {
            'data':
                [
                    entity
                    for entity in entity_list
                    if entity.get(key_field_name) == entity_id
                ] or None
        }
        for entity_id in entity_ids
    ]

    return results