import boto3
import io
from PIL import Image, ImageDraw, ExifTags, ImageColor, ImageFont
from supabase import create_client, Client

# Create an SSM client
ssm = boto3.client('ssm')

def get_parameter(name):
    params = {
        'Name': name,
        'WithDecryption': True
    }
    response = ssm.get_parameter(**params)
    return response['Parameter']['Value']

def lambda_handler(event, context):  
    try:
        # Fetch secrets from SSM Parameter Store
        SUPABASE_URL = get_parameter('/datafy-supabase-stack/supabase-url')
        SUPABASE_KEY = get_parameter('/datafy-supabase-stack/JwtSecret/ServiceRoleKey')
        REKOGNITION_MODEL_ARN = get_parameter('/datafy-rekognition-stack/model-arn')

        # Connect to Supabase
        supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)
    except Exception as e:
        return {
            'statusCode': 500,
            'body': str(e)
        }
    # connect to supabase here
    bucket='datafy-prod-detection' 
    photo='200ml juice.jpeg'
    model=REKOGNITION_MODEL_ARN
    min_confidence=50
    output_bucket = 'datafy-prod-detection-results'

    label_count = show_custom_labels(model,bucket,output_bucket,photo, min_confidence)
    print("Custom labels detected: " + str(label_count))
    result = {
        "started": label_count 
    }
    return result


def display_image(bucket,output_bucket,photo,response):
    # Load image from S3 bucket
    s3_connection = boto3.resource('s3')

    s3_object = s3_connection.Object(bucket,photo)
    s3_response = s3_object.get()

    stream = io.BytesIO(s3_response['Body'].read())
    image=Image.open(stream)

    # Create a new image with the drawn bounding boxes and labels
    draw_image = ImageDraw.Draw(image)

    # # Ready image to draw bounding boxes on it.
    imgWidth, imgHeight = image.size
    # draw = ImageDraw.Draw(image)

    # calculate and display bounding boxes for each detected custom label
    print('Detected custom labels for ' + photo)
    print('Custom Lambel : ' + str(response))
    for customLabel in response['CustomLabels']:
        print('Label ' + str(customLabel['Name']))
        print('Confidence ' + str(customLabel['Confidence']))
        
        if 'Geometry' in customLabel:
            box = customLabel['Geometry']['BoundingBox']
            left = imgWidth * box['Left']
            top = imgHeight * box['Top']
            width = imgWidth * box['Width']
            height = imgHeight * box['Height']

            fnt = ImageFont.load_default()
            draw_image.text((left, top), customLabel['Name'], fill='#00d400', font=fnt)
            # draw.text((left,top), customLabel['Name'], fill='#00d400', font=fnt)

            print('Left: ' + '{0:.0f}'.format(left))
            print('Top: ' + '{0:.0f}'.format(top))
            print('Label Width: ' + "{0:.0f}".format(width))
            print('Label Height: ' + "{0:.0f}".format(height))

            points = (
                (left,top),
                (left + width, top),
                (left + width, top + height),
                (left , top + height),
                (left, top))
            draw_image.line(points, fill='#00d400', width=5)
            # draw.line(points, fill='#00d400', width=5)

    # Save the image to the output bucket
    output_key = f"rekognition-{photo}"
    output_object = s3_connection.Object(output_bucket, output_key)
    image_bytes = io.BytesIO()
    image.save(image_bytes, format='JPEG')
    image_bytes.seek(0)
    output_object.put(Body=image_bytes, ContentType='image/jpeg')
    print(f"Image saved to {output_bucket}/{output_key}")    

def show_custom_labels(model,bucket,output_bucket,photo, min_confidence):
    client=boto3.client('rekognition')

    #Call DetectCustomLabels
    response = client.detect_custom_labels(Image={'S3Object': {'Bucket': bucket, 'Name': photo}},
        MinConfidence=min_confidence,
        ProjectVersionArn=model)

    # For object detection use case, uncomment below code to display image.
    display_image(bucket,output_bucket,photo,response)

    return len(response['CustomLabels'])