import json
import boto3
import logging
from typing import List
import logging
from typing import List
import os
import boto3

logger = logging.getLogger(__name__)


logger = logging.getLogger()
logger.setLevel(logging.INFO)
from botocore.exceptions import NoCredentialsError, PartialCredentialsError


S3_BUCKET_NAME = "suno-data-uploads"


def assume_role(role_arn, session_name):
    sts_client = boto3.client("sts")
    try:
        response = sts_client.assume_role(
            RoleArn=role_arn, RoleSessionName=session_name
        )
        credentials = response["Credentials"]
        return credentials
    except (NoCredentialsError, PartialCredentialsError) as e:
        print(f"Error assuming role: {e}")
        return None


def check_at_least_one_wav_exists(
    s3_client, bucket: str, clips_s3_data: List[List[str]]
) -> bool:
    """
    Check if '.wav' files exist for any of the given clips in the specified S3 bucket.

    This function iterates through the provided clips and checks for the existence
    of a '.wav' file corresponding to each clip in the S3 bucket. If it finds
    at least one such file, it returns True, indicating the presence of '.wav' files.

    Currently, helper function for delete_clip().

    Args:
        s3_client: The boto3 S3 client instance used for making API calls to AWS S3.
        bucket (str): The name of the S3 bucket where the files are stored.
        clips (Iterable): An iterable of clip objects, each expected to have an 's3_id'
                          attribute or an 'id' that can be used to construct the S3 key.

    Returns:
        bool: True if at least one '.wav' file is found for the given clips in the bucket,
              False otherwise.
    """
    for clip_s3_data in clips_s3_data:
        clip_id, s3_id, _ = clip_s3_data
        clip_s3_id = str(s3_id) if s3_id else str(clip_id)
        key = f"studio/uploads/{clip_s3_id}.wav"
        try:
            s3_client.head_object(Bucket=bucket, Key=key)
            return True
        except s3_client.exceptions.NoSuchKey:
            pass
        except s3_client.exceptions.ClientError as e:
            if e.response["Error"]["Code"] == "404":
                pass
            else:
                logger.error(f"Error in S3 operation: {str(e)}")
            return False

    return False


def move_image_file_to_deleted_if_exist(
    s3_client, bucket_name: str, file_id: str, file_ext: str
) -> None:
    """
    Move a image file to the 'deleted' folder within the S3 bucket.
    This function only applies to image files, it is backward compatible with the old image file format.

    Currently, helper function for delete_clip().

    Args:
        s3_client: The boto3 S3 client instance used for making API calls to AWS S3.
        bucket_name (str): The name of the S3 bucket.
        file_id (str): The identifier of the file (clip ID or image ID).
        file_ext (str): The file extension (e.g., 'mp3', 'png').
    """
    original_key = f"studio/uploads/{file_id}.{file_ext}"

    target_key = f"studio/deleted/{file_id}.{file_ext}"

    copy_source: "CopySourceTypeDef" = {"Bucket": bucket_name, "Key": original_key}

    try:
        s3_client.head_object(Bucket=bucket_name, Key=original_key)

        s3_client.copy(copy_source, bucket_name, target_key)
        s3_client.delete_object(Bucket=bucket_name, Key=original_key)
        logger.info(f"Moved {original_key}.{file_ext} to {target_key}.{file_ext}")
    except s3_client.exceptions.NoSuchKey:
        logger.info(
            f"File {original_key} does not exist, image file might be in different format."
        )
    except s3_client.exceptions.ClientError as e:
        if e.response["Error"]["Code"] == "404":
            logger.info(
                f"File {original_key} does not exist, image file might be in different format."
            )
        else:
            logger.error(f"Error in S3 operation: {str(e)}")


def move_file_to_deleted(
    s3_client, bucket_name: str, file_id: str, file_ext: str
) -> None:
    """
    Move a file to the 'deleted' folder within the S3 bucket. This can apply to audio files, images, etc.

    Currently, helper function for delete_clip().

    Args:
        s3_client: The boto3 S3 client instance used for making API calls to AWS S3.
        bucket_name (str): The name of the S3 bucket.
        file_id (str): The identifier of the file (clip ID or image ID).
        file_ext (str): The file extension (e.g., 'mp3', 'png').
    """
    original_key = f"studio/uploads/{file_id}.{file_ext}"

    target_key = f"studio/deleted/{file_id}.{file_ext}"

    copy_source = {"Bucket": bucket_name, "Key": original_key}

    try:
        s3_client.copy(copy_source, bucket_name, target_key)
        s3_client.delete_object(Bucket=bucket_name, Key=original_key)
    except s3_client.exceptions.NoSuchKey:
        pass
    except s3_client.exceptions.ClientError as e:
        if e.response["Error"]["Code"] == "404":
            pass
        else:
            logger.error(f"Error in S3 operation: {str(e)}")


def delete_file(s3_client, bucket_name: str, file_id: str, file_ext: str) -> None:
    original_key = f"studio/uploads/{file_id}.{file_ext}"
    try:
        s3_client.delete_object(Bucket=bucket_name, Key=original_key)
    except s3_client.exceptions.NoSuchKey:
        pass


def handler(event, context):
    # Create a Secrets Manager client
    client = boto3.client('secretsmanager')

    # Retrieve the secret value
    aws_creds = client.get_secret_value(SecretId='s3-bucket-creds')
    aws_creds = json.loads(aws_creds["SecretString"])

    aws_access_key_id = aws_creds['AWS_ACCESS_KEY_ID']
    aws_secret_access_key = aws_creds['AWS_SECRET_ACCESS_KEY']

    s3_client = boto3.client(
        "s3",
        aws_access_key_id=aws_access_key_id,
        aws_secret_access_key=aws_secret_access_key,
        region_name="us-east-1",
    )

    logging.info(f"Received event: {json.dumps(event)}")

    for record in event["Records"]:
        try:
            clips_s3_data = json.loads(record["body"])

            file_extensions = ["mp3", "npz", "m4a"]
            if check_at_least_one_wav_exists(s3_client, S3_BUCKET_NAME, clips_s3_data):
                file_extensions.append("wav")

            for clip_s3_data in clips_s3_data:
                clip_id, s3_id, image_s3_id = clip_s3_data
                clip_s3_id = str(s3_id) if s3_id else str(clip_id)

                # Handling audio files
                for ext in file_extensions:
                    move_file_to_deleted(s3_client, S3_BUCKET_NAME, clip_s3_id, ext)

                # Handling associated images
                if image_s3_id:
                    move_image_file_to_deleted_if_exist(
                        s3_client, S3_BUCKET_NAME, image_s3_id, "png"
                    )
                    move_image_file_to_deleted_if_exist(
                        s3_client,
                        S3_BUCKET_NAME,
                        image_s3_id.replace("image_", "image_large_"),
                        "png",
                    )
                    move_image_file_to_deleted_if_exist(
                        s3_client, S3_BUCKET_NAME, image_s3_id, "jpeg"
                    )
                    move_image_file_to_deleted_if_exist(
                        s3_client, S3_BUCKET_NAME, image_s3_id + "_small", "jpeg"
                    )

            # delete video file explicitly
            video_extensions = ["mp4"]
            for ext in video_extensions:
                delete_file(s3_client, S3_BUCKET_NAME, clip_s3_id, ext)

        except Exception as e:
            logger.error(f"Error processing SQS message: {str(e)}")
            continue

    return {
        "statusCode": 200,
        "body": json.dumps("Processed SQS messages successfully"),
    }