import xml.etree.ElementTree as ET
import json
import boto3
import os
from PIL import Image
import io
import logging
import time
import urllib.parse

# Set up logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

def lambda_handler(event, context):
    # Get the bucket and object from the event 
    records = event['Records'][0]
    source_bucket = records['s3']['bucket']['name'] 
    encoded_file_name = records['s3']['object']['key']
    print(f'XML file name  before decoding : {encoded_file_name}') 
    xml_file_decoded = urllib.parse.unquote(encoded_file_name)
    xml_file = xml_file_decoded.replace('+', ' ') 
    print(f'XML file name after decoding: {xml_file}')
    # Extract image file name (assuming it's in the same path as XML)
    image_file = os.path.splitext(xml_file)[0] + '.jpg'  # Adjust extension as needed
    original_filename = os.path.splitext(os.path.basename(image_file))[0]  # Get the base name without extension

    process_image_and_xml(source_bucket, xml_file, image_file, original_filename)

    return {
        'statusCode': 200,
        'body': json.dumps({'message': 'Images processed successfully'})
    }

def process_image_and_xml(source_bucket, xml_file, image_file, original_filename):
    # Initialize S3 client
    s3_client = boto3.client('s3')

    # Read XML content from S3
    xml_object = s3_client.get_object(Bucket=source_bucket, Key=xml_file)
    xml_content = xml_object['Body'].read()

    # Parse the XML content
    root = ET.fromstring(xml_content)

    # Retry mechanism for getting the image file
    max_attempts = 12  # 1 minute / 5 seconds
    attempt = 0

    while attempt < max_attempts:
        try:
            # Try to get the image from S3
            img_object = s3_client.get_object(Bucket=source_bucket, Key=image_file)
            img_content = img_object['Body'].read()
            break  # Exit loop if successful
        except s3_client.exceptions.NoSuchKey:
            logger.warning(f"Attempt {attempt + 1}: Image file {image_file} not found. Retrying in 5 seconds...")
            time.sleep(5)  # Wait for 5 seconds before retrying
            attempt += 1

    if attempt == max_attempts:
        logger.error(f"Could not retrieve image file {image_file} after {max_attempts} attempts.")
        return {
            'statusCode': 404,
            'body': json.dumps({'message': f'Could not get the file {image_file} after multiple attempts.'})
        }

    # Load image using PIL after successful retrieval
    image = Image.open(io.BytesIO(img_content))

    # Loop through each object in the XML to crop and save images
    for idx, obj in enumerate(root.findall('object')):
        label = obj.find('name').text  # Use the label as the folder name
        bndbox = obj.find('bndbox')
        
        # Extract bounding box coordinates
        xmin = int(bndbox.find('xmin').text)
        ymin = int(bndbox.find('ymin').text)
        xmax = int(bndbox.find('xmax').text)
        ymax = int(bndbox.find('ymax').text)

        logger.info(f"Processing object {idx + 1}: {label} with bbox ({xmin}, {ymin}, {xmax}, {ymax})")

        try:
            # Crop the image using bounding box coordinates
            cropped_image = image.crop((xmin, ymin, xmax, ymax))

            # Save cropped image to a BytesIO object
            output_buffer = io.BytesIO()
            cropped_image.save(output_buffer, format='JPEG')  # Adjust format as needed
            output_buffer.seek(0)

            # Define output file name based on label and original filename with a unique index
            output_folder_name = label  # Use label as folder name
            output_file_name = f"{output_folder_name}/{original_filename}_{idx + 1}.jpg"  # Unique filename

            # Upload cropped image to S3 in a different bucket 
            destination_bucket = 'datafy-prod-training'  # Destination bucket name
            logger.info(f"Saving cropped image to {output_file_name} in bucket {destination_bucket}")
            
            s3_client.put_object(Bucket=destination_bucket, Key=output_file_name, Body=output_buffer.getvalue())

        except Exception as e:
            logger.error(f"Error processing object {idx + 1}: {str(e)}")