import base64
import io
import logging

import boto3
from PIL import Image

# lambda function link: https://us-east-1.console.aws.amazon.com/lambda/home?region=us-east-1#/functions/image-resize?tab=code
#
# this lambda is to handle the image is not in jpeg format issue
#
# steps to update the functions:
# 1. updating the code
# 2. in the lambda code folder run the following command to download the PIL dependences
#   pip install \
#     --platform manylinux2014_x86_64 \
#     --only-binary=:all: --upgrade \
#     Pillow -t .
# 3. create a zip file with the following command in this folder
#   zip -r function.zip .
#
# 4. open the lambda and test the lambda is working
# 5. publish the version in the lambda
# 6. go to cloudfront and update the lamdbda in the origin response step


kinesisClient = boto3.client("kinesis")
cloudwatchClient = boto3.client("cloudwatch")
s3Client = boto3.client("s3")

# bucketName = 'sun-upload-testing-bucket'
# prefix = 'studio/upload/'

bucketName = "suno-data-uploads"
prefix = "studio/uploads/"

logger = logging.getLogger(__name__)


def lambda_handler(event, context):
    response = event["Records"][0]["cf"]["response"]
    request = event["Records"][0]["cf"]["request"]

    eventName = "bypass"

    try:
        respStatus = event["Records"][0]["cf"]["response"]["status"]
        fullFileName = request["uri"].split("/")[-1].lower()

        # if the request is png, will send an event to kinesis and another lambda will consume the event and create a jpeg iamge
        if fullFileName.endswith((".png")):
            kinesisClient.put_record(StreamName="image-resizing", Data="", PartitionKey=fullFileName)
            eventName = "send_to_kinesis"
        # if the request is jpeg and it missing in our s3, then load the png and create one jpeg and return to frontend
        elif fullFileName.endswith((".jpeg")) and (respStatus == "403" or respStatus == "404"):
            filename = fullFileName[:-5]
            png = prefix + filename + ".png"
            try:
                s3Object = s3Client.get_object(Bucket=bucketName, Key=png)
                s3Data = s3Object["Body"].read()
                image = Image.open(io.BytesIO(s3Data))
                rgbIm = image.convert("RGB")
                buffer = io.BytesIO()
                rgbIm.save(buffer, format="JPEG")
                buffer.seek(0)
                imgBytes = buffer.read()
                s3Client.put_object(
                    Bucket=bucketName, Key=prefix + fullFileName, Body=imgBytes, ContentType="image/jpeg"
                )
                eventName = "success_create"
                response = {
                    "status": "200",
                    "statusDescription": "OK",
                    "headers": {
                        "content-type": [{"key": "Content-Type", "value": "image/jpeg"}],
                        "transfer-encoding": [{"key": "Transfer-Encoding", "value": "chunked"}],
                    },
                    "body": base64.b64encode(imgBytes).decode("utf-8"),
                    "bodyEncoding": "base64",
                }

            except Exception as e:
                logger.error(e)
                eventName = "failed_gen_jpeg"

    except Exception as e:
        logger.error(e)
        eventName = "unknown_error"

    # recording the event to the metrics
    cloudwatchClient.put_metric_data(
        Namespace="Cloudfront",
        MetricData=[
            {
                "MetricName": "ImageResizer",
                "Dimensions": [{"Name": "EventName", "Value": eventName}],
                "Value": 1,
                "Unit": "Count",
            },
        ],
    )

    return response