import replicate
from pathlib import Path
import uuid
import requests
from enum import Enum

class Styles(Enum):
    emoji = "emoji"
    pixel_art = "pixel art"
    svg = "svg"
    cartoon = "cartoon"


def image_style(base_prompt: str, style: Styles) -> str:
    return f"create {style.name} style image on a white background: {base_prompt}"


def replicate_flux_schnell_image_gen(prompt: str) -> str:
    output = replicate.run(
        "black-forest-labs/flux-schnell",
        input={
            "prompt": prompt,
            "go_fast": True,
            "megapixels": "1",
            "num_outputs": 1,
            "aspect_ratio": "1:1",
            "output_format": "png",
            "output_quality": 80,
            "num_inference_steps": 4
        }
    )
    return str(output[0])


def save_image_url(image_url: str, output_dir: Path):
    output_dir.mkdir(parents=True, exist_ok=True)
    image_filename = str(uuid.uuid4()) + ".png"
    print(image_filename)
    image_filepath = output_dir / image_filename

    # Download the image
    image_response = requests.get(image_url)
    image_response.raise_for_status()
    with open(image_filepath, "wb") as f:
        f.write(image_response.content)


def replicate_flux_1_1_pro_image_gen(prompt: str) -> str:
    output = replicate.run(
        "black-forest-labs/flux-1.1-pro",
        input={
            "prompt": prompt,
            "aspect_ratio": "1:1",
            "output_format": "png",
            "output_quality": 100,
            "safety_tolerance": 1,
            "prompt_upsampling": True
        }
    )
    return str(output)


if __name__ == "__main__":
    base_prompt = "A playful, cartoon-style dog icon featuring a round, fluffy body in a warm golden-brown color. The dog has oversized, expressive black eyes and a happy, open mouth showing its tongue out. Two large floppy ears hang down on either side of its head, each with a slightly darker brown color on the inside. The tail is wagging and curved upwards, represented by a simple, thick shape in the same golden-brown. The background is a light blue with a subtle oval shape to make the dog stand out boldly. This design uses strong shapes and clear contrasts for easy recognition in a small format."
    dog_emoji = image_style(base_prompt, style=Styles.cartoon)
#    image_url = replicate_flux_schnell_image_gen(dog_emoji)
#    output_dir = Path("../output/test2/replicate/flux_schnell")
    image_url = replicate_flux_1_1_pro_image_gen(dog_emoji)
    output_dir = Path("../output/test2/replicate/flux_1_1_pro")
    save_image_url(image_url, output_dir)