import boto3
import os
import zipfile
import io
import logging

# Configure logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# Initialize S3 client
s3_client = boto3.client("s3")


def lambda_handler(event, context):
    """
    AWS Lambda function that is triggered when a file is uploaded to an S3 bucket.
    If the file is a zip file, it extracts the contents and saves them to the same bucket,
    then deletes the original zip file.

    Args:
        event (dict): Event data from S3 trigger
        context (object): Lambda context

    Returns:
        dict: Response with processing results
    """
    try:
        # Get bucket and key from the S3 event
        bucket = event["Records"][0]["s3"]["bucket"]["name"]
        key = event["Records"][0]["s3"]["object"]["key"]

        logger.info(f"Processing file: {key} from bucket: {bucket}")

        # Check if the file is a zip file
        if not key.lower().endswith(".zip"):
            logger.info(f"File {key} is not a zip file. Skipping.")
            return {
                "statusCode": 200,
                "body": f"File {key} is not a zip file. No action taken.",
            }

        # Get the zip file content
        response = s3_client.get_object(Bucket=bucket, Key=key)
        zip_content = response["Body"].read()

        # Process the zip file
        with zipfile.ZipFile(io.BytesIO(zip_content)) as zip_ref:
            file_list = zip_ref.namelist()
            logger.info(f"Found {len(file_list)} files in the zip archive")

            # Extract each file and upload to S3
            for file_path in file_list:
                # Skip directories
                if file_path.endswith("/"):
                    continue

                # Extract file content
                file_content = zip_ref.read(file_path)

                # Determine content type (optional)
                content_type = "application/octet-stream"  # Default
                if file_path.lower().endswith(".txt"):
                    content_type = "text/plain"
                elif file_path.lower().endswith(".xml"):
                    content_type = "text/xml"
                elif file_path.lower().endswith(".jpg") or file_path.lower().endswith(
                    ".jpeg"
                ):
                    content_type = "image/jpeg"
                elif file_path.lower().endswith(".png"):
                    content_type = "image/png"
                elif file_path.lower().endswith(".avif"):
                    content_type = "image/avif"
                elif file_path.lower().endswith(".webp"):
                    content_type = "image/webp"

                # Create the new key for the extracted file
                # Remove the .zip extension from the original key and use as prefix
                prefix = os.path.splitext(key)[0]
                new_key = f"{prefix}/{file_path}"

                # Upload the file to S3
                s3_client.put_object(
                    Bucket=bucket,
                    Key=new_key,
                    Body=file_content,
                    ContentType=content_type,
                )
                logger.info(f"Uploaded extracted file: {new_key}")

        # Delete the original zip file
        s3_client.delete_object(Bucket=bucket, Key=key)
        logger.info(f"Deleted original zip file: {key}")

        return {
            "statusCode": 200,
            "body": f"Successfully processed zip file {key}. Extracted {len(file_list)} files.",
        }

    except Exception as e:
        logger.error(f"Error processing zip file: {str(e)}")
        return {"statusCode": 500, "body": f"Error processing zip file: {str(e)}"}