import os
from typing import List
import replicate
import base64
from dotenv import load_dotenv


load_dotenv()
replicate_client = replicate.Client(api_token= os.getenv("REPLICATE_API_KEY"))

def image_generation_change_background(
    img_bytes: bytes,
    prompt: str,
    aspect_ratio: str):
    encoded_image = base64.b64encode(img_bytes).decode("utf-8")

    output = replicate_client.run(
        "google/nano-banana",
        input={
            "image_input": [f"data:image/jpeg;base64,{encoded_image}"],
            "prompt": prompt,
            "aspect_ratio": aspect_ratio,
        }
    )
    urls: List[str] = []
    if isinstance(output, list) and output:
        first = output[0]
        url = getattr(first, "url", str(first))
        urls = [url]
    elif isinstance(output, str):
        urls = [output]
    elif hasattr(output, "url"):
        urls = [getattr(output, "url")]
    if urls:
        return urls[0]