import json
import modal
import uuid


def generate_image(prompt_text: str) -> str:
    """
    Generate an image using the Flux StableDiffusion model.

    Args:
        prompt_text (str): The text prompt describing the image to generate

    Returns:
        str: The UUID of the generated image
    """

    image_id = str(uuid.uuid4())

    model_f = modal.Cls.from_name(f"flux-dev", "StableDiffusion")
    model_f().generate_image_item.spawn(
        json.dumps(
            {
                "id": image_id,
                "prompt_text": prompt_text,
                "metadata": {},
            },
        ),
    )

    return image_id


def download_image_from_s3(image_id: str):
    """
    Retrieve an image from S3 using its ID.

    Args:
        image_id (str): The UUID of the image to retrieve

    Returns:
        str: The local path to the downloaded image
    """
    import os
    import subprocess

    # Create the S3 path for the image
    s3_path = f"s3://suno-data-uploads/studio/uploads/image_large_{image_id}.jpeg"

    # Define local path for the downloaded image
    local_path = f"image_{image_id}.jpeg"

    # Use AWS CLI to download the file
    try:
        subprocess.run(["aws", "s3", "cp", s3_path, local_path], check=True)
        return local_path
    except subprocess.CalledProcessError as e:
        print(f"Error downloading image from S3: {e}")
        return None