import json
import boto3
import io
import gzip
import base64
from io import StringIO

mf = io.BytesIO()
gg = boto3.client('greengrass')
iot = boto3.client('iot')

def gzip_b64encode(data):
    compressed = io.BytesIO()
    with gzip.GzipFile(fileobj=compressed, mode='w') as f:
        json_response = json.dumps(data)
        f.write(json_response.encode('utf-8'))
    return base64.b64encode(compressed.getvalue()).decode('ascii')

def lambda_handler(event, context):

    group = gg.create_group(Name="group_test")

    keys_cert = iot.create_keys_and_certificate(setAsActive=True)
    core_thing = iot.create_thing(thingName="core_group_test")
    keys = keys_cert['keyPair']

    cert_file = io.StringIO(keys_cert['certificatePem'])

    private_key_file = io.StringIO(keys['PrivateKey'])

    public_key_file = io.StringIO(keys['PublicKey'])

    iot.attach_thing_principal(
        thingName=core_thing['thingName'],
        principal=keys_cert['certificateArn'])

    core_policy_doc = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Action": ["iot:Publish", "iot:Subscribe", "iot:Connect", "iot:Receive", "iot:GetThingShadow", "iot:DeleteThingShadow", "iot:UpdateThingShadow"],
                "Resource": ["arn:aws:iot:" + boto3.session.Session().region_name + ":*:*"]
            },
            {
                "Effect": "Allow",
                "Action": ["greengrass:AssumeRoleForGroup", "greengrass:CreateCertificate", "greengrass:GetConnectivityInfo", "greengrass:GetDeployment", "greengrass:GetDeploymentArtifacts", "greengrass:UpdateConnectivityInfo", "greengrass:UpdateCoreDeploymentStatus"],
                "Resource": ["*"]
            }
        ]
    }
    policy = iot.create_policy(
        policyName="my_core_23_policy",
        policyDocument=json.dumps(core_policy_doc))

    iot.attach_principal_policy(
        policyName=policy['policyName'],
        principal=keys_cert['certificateArn'])

    initial_version = {'Cores': [
        {
            'Id': core_thing['thingName'], # Quite intuitive, eh?
            'CertificateArn': keys_cert['certificateArn'],
            'SyncShadow': False, # Up to you, True|False
            'ThingArn': core_thing['thingArn']
        }
    ]}

    core_definition = gg.create_core_definition(
        Name="{0}_core_def".format(group['Name']),
        InitialVersion=initial_version)

    retrieved_data = {}
    retrieved_data['cert.pem'] = keys_cert['certificatePem']
    retrieved_data['private.key'] = keys['PrivateKey']
    retrieved_data['public.key'] = keys['PublicKey']
    # # Create in memory zip and add files
    # zippedfiles = zipfile.ZipFile(mf, mode='w',compression=zipfile.ZIP_DEFLATED)
    # zippedfiles.writestr('cert.pem', io.StringIO(keys_cert['certificatePem'])
    # zippedfiles.writestr('private.key', io.StringIO(keys_cert['certificatePem']['PrivateKey'])
    # zippedfiles.writestr('public.key', io.StringIO(keys_cert['certificatePem']['PublicKey'])
    # zippedfiles.close()

    # with zipfile.ZipFile(mf, mode="w",compression=zipfile.ZIP_DEFLATED) as zf:
    #     zf.writestr('cert.pem', io.StringIO(keys_cert['certificatePem']))
    #     # zf.writestr('private.key', io.StringIO(keys_cert['certificatePem']['PrivateKey'])
    #     # zf.writestr('public.key', io.StringIO(keys_cert['certificatePem']['PublicKey'])

    # with open("zf.zip", "wb") as f: # use `wb` mode
    #     f.write(mf.getvalue())

    return {
            "isBase64Encoded": True,
            "statusCode": 200,
            "headers": {
                "Content-Type": "application/zip",
                #"Content-Encoding": "gzip",
                #"Content-Disposition": "attachment; filename=test.zip",
                #"Access-Control-Allow-Origin": "*"
            },
            "body": gzip_b64encode(retrieved_data)
        }