import boto3
import io
import requests  # Import requests to fetch images from URLs
from PIL import Image, ImageDraw, ImageFont
from supabase import create_client, Client

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

# Initialize global variables
SUPABASE_URL = None
SUPABASE_KEY = None
supabase: Client = None
REKOGNITION_MODEL_ARN = None

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

def initialize_parameters():
    global SUPABASE_URL, SUPABASE_KEY, supabase, REKOGNITION_MODEL_ARN
    if SUPABASE_URL is None or SUPABASE_KEY is None or REKOGNITION_MODEL_ARN is None:
        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')
        supabase = create_client(SUPABASE_URL, SUPABASE_KEY)

def lambda_handler(event, context):  
    try:
        # Initialize parameters and Supabase client
        initialize_parameters()

    except Exception as e:
        return {
            'statusCode': 500,
            'body': str(e)
        }

    # Query product_images table for id, product_id, and image_path where processed is false
    try:
        response = supabase.table('product_images').select('id, product_id, image_path').eq('processed', False).execute()
        
        if not response.data:
            return {
                'statusCode': 404,
                'body': "No unprocessed images found in the product_images table."
            }

        # Iterate through the results and process each image
        for item in response.data:
            record_id = item['id']  # Get the unique identifier for updating later
            product_id = item['product_id']
            image_path = item['image_path']  # This is the URL to the image

            # Call show_custom_labels with the image URL instead of S3 bucket details
            rekognition_response = show_custom_labels(image_path)
            print(f"Custom labels detected for product {product_id}: {rekognition_response}")

            # Save the new image and results to Supabase Storage and product_ml table
            new_image_url = save_image_to_supabase(rekognition_response, image_path) 

            # Insert data into product_ml table
            supabase.table('product_ml').insert({
                'product_id': product_id,
                'ml_image_path': new_image_url,
                'ml_result': rekognition_response,
                'product_images_id': record_id
            }).execute()

            # Update the processed field in Supabase after processing
            supabase.table('product_images').update({'processed': True}).eq('id', record_id).execute()

    except Exception as e:
        print("Error fetching data from Supabase:", str(e))
        return {
            'statusCode': 500,
            'body': str(e)
        }

    return {
        'statusCode': 200,
        'body': "Processed images successfully."
    }

def display_image(image_url, response):
    # Load image from the URL instead of S3
    try:
        img_response = requests.get(image_url)
        img_response.raise_for_status()  # Raise an error for bad responses

        stream = io.BytesIO(img_response.content)
        image = Image.open(stream)

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

        imgWidth, imgHeight = image.size

        print('Detected custom labels for ' + image_url)
        
        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)

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

        return image  # Return modified image for saving

    except Exception as e:
        print("Error loading or processing image:", str(e))

def show_custom_labels( image_url):
    client=boto3.client('rekognition')

    try:
        # Fetching the image from URL and converting it to bytes
        img_response = requests.get(image_url)
        img_response.raise_for_status()  # Raise an error for bad responses
        
        # Convert the image to bytes
        img_bytes = io.BytesIO(img_response.content)

        # Call DetectCustomLabels using base64-encoded bytes of the image
        response = client.detect_custom_labels(
            Image={'Bytes': img_bytes.getvalue()},
            MinConfidence=50,
            ProjectVersionArn=REKOGNITION_MODEL_ARN
        )

    except Exception as e:
        print("Error detecting custom labels:", str(e))
        return []

    return response

def save_image_to_supabase(rekognition_response, image_path):
    try:
        modified_image = display_image(image_path,rekognition_response)  # Assuming display_image returns modified PIL Image
        
        output_stream = io.BytesIO()
        
        # Convert the image to RGB if it's in RGBA mode
        if modified_image.mode == 'RGBA':
            modified_image = modified_image.convert('RGB')
            
        modified_image.save(output_stream, format='JPEG')  # Save modified image to bytes stream
        
        output_stream.seek(0)  # Reset stream position
        
        # Extracting filename from original image URL using split
        parts = image_path.split('/')  # Split by '/'
        original_filename = parts[-1]  # Get the last part, which is the filename
        file_name = f"product/{original_filename.split('.')[0]}.jpg"  
        
        upload_response = supabase.storage.from_('ml').upload(file_name, output_stream.getvalue(), {
            "contentType": "image/jpeg"
        })

        print(f'Upload response : {upload_response}')
        # Get public URL of uploaded file
        public_url_data = supabase.storage.from_('ml').get_public_url(file_name)

        print(f'Public url response : {public_url_data}')
        
        return public_url_data.rstrip('?')

    except Exception as e:
        print("Error saving image to Supabase:", str(e))