import gradio as gr
from PIL import Image, ImageEnhance, ImageOps
import numpy as np

# Styling presets for prompt expansion
STYLES = {
    "None (Pass-through)": "{prompt}",
    "Photorealistic": "{prompt}, macro photography, extremely shallow depth of field, sharp focus, natural lighting, high-contrast minimal composition, warm skin tones, cinematic color palette, shot on 85mm lens, 8k resolution, photorealistic",
    "3D Toy Figure": "3D rendered matte {prompt} toy figure, stylized round anthropomorphic shape, smooth vinyl texture, studio lighting, solid vibrant background, high contrast, minimal composition, octane render, raytracing, 3d art",
    "Anime / Manga": "highly detailed digital painting of {prompt}, anime key art style, vibrant color palette, dynamic lighting, beautiful eyes, dramatic angle, concept art aesthetic, studio ghibli or makoto shinkai style",
    "Cyberpunk": "retro-futuristic cyberpunk style {prompt}, neon glow, wet streets with reflections, holographic details, cinematic lighting, dark atmosphere, blade runner aesthetic, highly detailed, 8k",
    "Ligne Claire / Minimalist": "minimalist flat-color illustration of {prompt}, clean lines, delicate paper texture, vast negative space, high-angle perspective, harmonious color palette, ligne claire style, modern vector graphic",
    "Surreal Oil Painting": "surreal dreamlike painting of {prompt}, thick impasto oil paint texture, visible coarse brushstrokes, rich color blending, atmospheric lighting, masterpiece, gallery quality"
}

def enhance_prompt(prompt: str = "immense rocket launch exhaust as seen from extremely close up", style: str = "None (Pass-through)") -> str:
    """
    Enhance a base prompt with descriptive style templates.
    """
    if not prompt:
        return "immense rocket launch exhaust as seen from extremely close up"
    
    template = STYLES.get(style, "{prompt}")
    return template.replace("{prompt}", prompt)

def apply_filter(image: Image.Image, filter_type: str = "None") -> Image.Image:
    """
    Apply a PIL-based aesthetic filter to the input image.
    """
    if image is None:
        return None
        
    # Ensure it is a PIL Image
    if not isinstance(image, Image.Image):
        try:
            image = Image.fromarray(np.array(image).astype('uint8'), 'RGB')
        except Exception:
            return image
            
    # Convert to RGB if not already
    if image.mode != "RGB":
        image = image.convert("RGB")
        
    if filter_type == "None":
        return image
    elif filter_type == "Black & White":
        return ImageOps.grayscale(image)
    elif filter_type == "Sepia / Vintage":
        # Vintage sepia filter
        width, height = image.size
        pixels = image.load()
        for py in range(height):
            for px in range(width):
                r, g, b = pixels[px, py]
                tr = int(0.393 * r + 0.769 * g + 0.189 * b)
                tg = int(0.349 * r + 0.686 * g + 0.168 * b)
                tb = int(0.272 * r + 0.534 * g + 0.131 * b)
                pixels[px, py] = (min(tr, 255), min(tg, 255), min(tb, 255))
        return image
    elif filter_type == "High Contrast":
        enhancer = ImageEnhance.Contrast(image)
        return enhancer.enhance(1.6)
    elif filter_type == "Cool / Cyberpunk":
        # Enhance blue/cyan, lower red
        r, g, b = image.split()
        r = r.point(lambda i: i * 0.8)
        b = b.point(lambda i: min(255, int(i * 1.3)))
        return Image.merge("RGB", (r, g, b))
    elif filter_type == "Warm / Golden Hour":
        # Enhance red/yellow
        r, g, b = image.split()
        r = r.point(lambda i: min(255, int(i * 1.25)))
        g = g.point(lambda i: min(255, int(i * 1.1)))
        b = b.point(lambda i: i * 0.8)
        return Image.merge("RGB", (r, g, b))
        
    return image

# Initialize the Gradio Workflow, binding our custom nodes
gr.Workflow(
    graph="workflow.json",
    bind={
        "Enhance Prompt": enhance_prompt,
        "Apply Filter": apply_filter
    }
).launch()