import json
import boto3
import uuid
from time import time
import cv2

# XRay instrumentation: https://docs.aws.amazon.com/lambda/latest/dg/python-tracing.html
from aws_xray_sdk.core import patch_all
patch_all()

s3_client = boto3.client('s3')

tmp = "/tmp/"
FILE_NAME_INDEX = 0
FILE_PATH_INDEX = 2


def video_processing(object_key, video_path):
    file_name = object_key.split(".")[FILE_NAME_INDEX]
    result_file_path = tmp+file_name+'-output.avi'

    video = cv2.VideoCapture(video_path)

    width = int(video.get(3))
    height = int(video.get(4))

    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter(result_file_path, fourcc, 20.0, (width, height))

    start = time()
    while video.isOpened():
        ret, frame = video.read()

        if ret:
            gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            tmp_file_path = tmp+'tmp.jpg'
            cv2.imwrite(tmp_file_path, gray_frame)
            gray_frame = cv2.imread(tmp_file_path)
            out.write(gray_frame)
        else:
            break

    latency = time() - start

    video.release()
    out.release()
    return latency, result_file_path


def lambda_handler(event, context):
    # NOTE: Asynchronous invocation does not require JSON parsing
    event_body = json.loads(event['body'])
    input_bucket = event_body['input_bucket']
    object_key = event_body['object_key']
    output_bucket = event_body['output_bucket']

    download_path = tmp+'{}{}'.format(uuid.uuid4(), object_key)

    s3_client.download_file(input_bucket, object_key, download_path)

    latency, upload_path = video_processing(object_key, download_path)

    s3_client.upload_file(upload_path, output_bucket, upload_path.split("/")[FILE_PATH_INDEX])

    body = {
        'latency': latency
    }
    response = {
        'statusCode': 200,
        'body': json.dumps(body)
    }
    return response
