import gradio as gr
import torch
import random
import time
from diffusers import DiffusionPipeline, LCMScheduler
from PIL import Image, ImageFilter

# -------------------------------
# MODEL SETUP (CPU SAFE)
# -------------------------------
MODEL_ID = "runwayml/stable-diffusion-v1-5"
ADAPTER_ID = "latent-consistency/lcm-lora-sdv1-5"

pipe = DiffusionPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,
    safety_checker=None
)
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.load_lora_weights(ADAPTER_ID)
pipe.to("cpu")

pipe.enable_attention_slicing()
pipe.enable_vae_slicing()
pipe.set_progress_bar_config(disable=True)

# -------------------------------
# FAST THINKING PROMPT ENGINE
# -------------------------------
def refine_prompt_fast(user_prompt: str):
    """
    Instant deterministic prompt refinement (<1ms)
    """
    known_objects = {"apple","banana","snake","cat","dog","fox","rabbit","dragon","bird","frog","hamster"}
    object_found = next((w for w in known_objects if w in user_prompt.lower()), None)
    subject = object_found if object_found else user_prompt.lower().strip()

    style = ""
    if any(w in user_prompt.lower() for w in ["cute","adorable","kawaii"]):
        style = ", cute, friendly, rounded body, big eyes, pastel colors, cartoon style"
    elif any(w in user_prompt.lower() for w in ["realistic","photo","photograph"]):
        style = ", ultra realistic, natural anatomy, professional photography"
    else:
        style = ", high quality, clean background"

    prompt = f"a single {subject}, centered, isolated{style}"
    negative = "multiple objects, duplicate, blurry, low quality, cropped, out of frame, horror, grotesque, aggressive, weird colors, artifacts"
    return prompt, negative

# -------------------------------
# REALISTIC ETA CALCULATION
# -------------------------------
def estimate_time(steps, resolution):
    per_step = {256:6, 512:12, 768:25, 1024:45}[int(resolution)]
    overhead = 2
    est = overhead + steps * per_step
    minutes = est // 60
    seconds = est % 60
    return f"⏱️ Estimated time: ~{int(minutes)}m {int(seconds)}s"

# -------------------------------
# IMAGE GENERATION WITH PROGRESSIVE BLUR GALLERY
# -------------------------------
def generate(prompt, resolution, steps):
    # --- THINKING PHASE ---
    refined_prompt, neg_prompt = refine_prompt_fast(prompt)

    # --- SHOW WHITE IMAGE WHILE GENERATING ---
    width, height = int(resolution), int(resolution)
    blank_img = Image.new("RGB", (width, height), (255, 255, 255))
    yield [blank_img], ""  # always show gallery placeholder

    # --- GENERATION PHASE ---
    seed = random.randint(0, 10**9)
    gen = torch.Generator("cpu").manual_seed(seed)
    pipe.scheduler.set_timesteps(int(steps))
    img = pipe(
        prompt=refined_prompt,
        negative_prompt=neg_prompt,
        num_inference_steps=int(steps),
        guidance_scale=1.2,
        width=width,
        height=height,
        generator=gen
    ).images[0]

    # --- PROGRESSIVE BLUR REVEAL ---
    max_blur = 20  # max blur radius for 100%
    steps_blur = 10
    for i in range(steps_blur):
        blur_percent = 100 - i*10
        blurred_img = img.filter(ImageFilter.GaussianBlur(radius=max_blur * blur_percent / 100))
        yield [blurred_img], ""  # gallery updated, no status text
        time.sleep(1)  # 1 second per step

    # --- FINAL IMAGE ---
    yield [img], ""  # fully revealed image

# -------------------------------
# GRADIO UI
# -------------------------------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 👾 CREEPER AI — CPU SMART IMAGE GENERATION")

    with gr.Row():
        with gr.Column():
            prompt_in = gr.Textbox(label="Prompt", placeholder="cute snake", lines=2)
            
            resolution = gr.Radio([256, 512, 768, 1024], value=512, label="Resolution")
            
            steps = gr.Slider(2, 8, value=4, step=1, label="Steps")
            
            eta = gr.Markdown("⏱️ Estimated time: ~1m 0s")
            gen_btn = gr.Button("Generate")
        
        with gr.Column():
            gallery = gr.Gallery(columns=1)

    # Update ETA dynamically
    for ctrl in [steps, resolution]:
        ctrl.change(estimate_time, [steps, resolution], eta)

    # Button click triggers generation
    gen_btn.click(
        generate,
        inputs=[prompt_in, resolution, steps],
        outputs=[gallery, gallery]  # gallery updated for progressive blur
    )

demo.launch()