import gzip
import json
import base64
import os 
import boto3

CENTRAL_ACCOUNT_ROLE= os.environ["CENTRAL_ACCOUNT_ROLE"]

def load_az(az,central_account_role):
    table={}
    if central_account_role == 'NoValue':
        ddb = boto3.resource('dynamodb')
        table = ddb.Table('AZsMapping')

    else:
        sts_connection = boto3.client('sts')
        acct_b = sts_connection.assume_role(
            RoleArn=central_account_role,
            RoleSessionName="cross_acct_lambda"
        )
        
        ACCESS_KEY = acct_b['Credentials']['AccessKeyId']
        SECRET_KEY = acct_b['Credentials']['SecretAccessKey']
        SESSION_TOKEN = acct_b['Credentials']['SessionToken']

        ddb = boto3.resource(
            'dynamodb',
            aws_access_key_id=ACCESS_KEY,
            aws_secret_access_key=SECRET_KEY,
            aws_session_token=SESSION_TOKEN,
        )
        table = ddb.Table('AZsMapping')

    table.put_item(Item=az)


def lambda_handler(event, context):
    # TODO implement
    print(f'Logging Event: {event}')
    print(f"Awslog: {event['awslogs']}")
    cw_data = event['awslogs']['data']
    print(f'data: {cw_data}')
    print(f'type: {type(cw_data)}')
    compressed_payload = base64.b64decode(cw_data)
    uncompressed_payload = gzip.decompress(compressed_payload)
    payload = json.loads(uncompressed_payload)
    az = {}
    
    log_events = payload['logEvents']
    response_element_json = json.loads(log_events[0]['message'])
    az['AvailabilityZoneId'] = response_element_json['responseElements']['subnet']['availabilityZoneId']
    az['CidrBlock'] = response_element_json['responseElements']['subnet']['cidrBlock']
    az['SubnetId'] = response_element_json['responseElements']['subnet']['subnetId']
    print(az)
    load_az(az,CENTRAL_ACCOUNT_ROLE)