import base64
import json
from datetime import datetime, timezone


def lambda_handler(event, context):
    """
    Lambda function to process rec-events data for Firehose delivery to S3.
    Adds partitioning keys and processes recommendation event data.
    """
    output = {'records': []}
    print(f"Processing {len(event['records'])} rec-events records")
    
    for record in event['records']:
        try:
            arrival_time = datetime.fromtimestamp(record['approximateArrivalTimestamp'] / 1000)
            
            # Decode and parse the payload for potential custom processing
            payload = base64.b64decode(record['data'])
            json_data = json.loads(payload)
            
            # Add processed timestamp to the data
            json_data['processed_timestamp'] = datetime.utcnow().isoformat() + 'Z'
            json_data['processing_source'] = 'rec-events-logger-parser'
            
            # Create partition keys for S3 organization
            partition_keys = {
                "EventVersion": 1,  # Version 1 for rec-events
                "year": str(arrival_time.year),
                "month": str(arrival_time.month).zfill(2),
                "date": str(arrival_time.day).zfill(2),
                "hour": str(arrival_time.hour).zfill(2),
            }
            
            # Re-encode the enriched data
            enriched_data = json.dumps(json_data) + '\n'
            encoded_data = base64.b64encode(enriched_data.encode('utf-8')).decode('utf-8')
            
            firehose_record_output = {
                'recordId': record['recordId'],
                'data': encoded_data,
                'result': 'Ok',
                'metadata': { 'partitionKeys': partition_keys }                
            }
            
        except Exception as e:
            print(f"Error processing record {record['recordId']}: {str(e)}")
            # Return the original record with ProcessingFailed result
            firehose_record_output = {
                'recordId': record['recordId'],
                'data': record['data'],
                'result': 'ProcessingFailed'
            }
        
        output['records'].append(firehose_record_output)
    
    print(f"Successfully processed {len(output['records'])} records")
    return output