import base64
import json


def firehose_response_formatter(recordId, processed_payload):
    """
    helper function that formats the processed records according to Firehose object response model.

    Payload is first encoded to convert it to bytes string. Then it is encoded with base64 since Firehose 
    requires a base64 encoded object in response. This base64 encoded string is further decoded to convert it
    to string object which is json serilizable and retunable by lambda.

    :param record: the recordId of the record
    :param processed_payload: the processed data
    :return: dict object 
    """

    # status could be Ok, Dropped, ProcessingFailed. Chose Ok if records are processed successfully
    # choose Dropped if a record is to be dropped.
    # ProcessingFailed auto selected in case of failure.

    output_record = {
        'recordId': recordId,
        'result': 'Ok',
        'data': base64.b64encode(json.dumps(processed_payload).encode('utf-8')).decode('utf-8')
    }
    return output_record


def lambda_handler(event, context):
    """
    lambda handler reads records put on the kinesis data firehose. Records are automatically deaggregated, hence
    only have to perform custom processing on the payload. The processed records are sent back to Firehose
    in the require object response model.

    records are recieved as base64 encoded string. Decoding them returns a bytes string 
    which is further decoded to get the actual string object

    :param event:  the event object generated by the service invoked
    :param context: This object provides methods and properties that provide information about the invocation, function\
    and execution environment.
    :return: dict object containing records
    """

    output = []

    for record in event['records']:
        recordId = record['recordId']
        payload = base64.b64decode(record['data']).decode('utf-8')

        # custom processing on payload here

        processed_payload = payload

        output.append(firehose_response_formatter(recordId, processed_payload))

    return {'records': output}
