import os
import requests
from PIL import Image
from io import BytesIO

API_KEY = "sk-vvj3pYatFrBEendEmpUKcX77IxndwUG3wXpix6cXVc9btaAE"

URL = "https://api.stability.ai/v2beta/stable-image/edit/replace-background-and-relight"


def edit_image(image_path: str, prompt: str):

    if API_KEY is None:
        raise Exception("Missing STABILITY_API_KEY")

    if not os.path.exists(image_path):
        raise Exception("Image not found")

    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Accept": "image/*"
    }

    with open(image_path, "rb") as img:

        files = {
            "subject_image": img
        }

        data = {

            # New background
            "background_prompt": prompt,

            # Preserve person
            "preserve_original_subject": 1,

            "output_format": "png"

        }

        response = requests.post(
            URL,
            headers=headers,
            files=files,
            data=data,
            timeout=300
        )

    print(response.status_code)

    if response.status_code != 200:
        print(response.text)
        raise Exception(response.text)

    return Image.open(BytesIO(response.content)).convert("RGB")