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

def create_image(filename, prompt):
    print(f"Creating image: {filename}")
    prompt = re.sub(r'\s+', ' ', prompt).strip()

    url = f"https://pollinations.ai/p/{prompt}?seed=2025&nologo=true&model=flux&width=1080&height=1080"
    response = requests.get(url)
    original_image = Image.open(BytesIO(response.content)).convert("RGB")

    base_filename = os.path.splitext(filename)[0]
    save_path = os.path.join("agent", f"{base_filename}_1080.png")
    os.makedirs("agent", exist_ok=True)
    original_image.save(save_path)
    print(f"Saved original image: {save_path}")

    original_path = os.path.join("agent", filename)
    original_image.save(original_path)
    print(f"Saved original image as: {original_path}")

    sizes = {
        "512": (512, 512),
        "192": (192, 192),
        "48": (48, 48),
    }

    for size_label, size in sizes.items():
        resized = original_image.resize(size, Image.LANCZOS)
        resized_filename = os.path.join("agent", f"{base_filename}_{size_label}.png")
        resized.save(resized_filename)
        print(f"Saved resized image: {resized_filename}")

if __name__ == "__main__":
    filename = "example_image.png"  # Default filename, can be changed
    prompt = "A beautiful landscape with mountains and a river"  # Default prompt, can be changed
    
    if filename and prompt:
        create_image(filename, prompt)
    else:
        print("Filename and prompt cannot be empty.")