# app.py — MangaMorph (Gradio) — backwards-compatible, CPU-friendly
import os
import random
import time
import numpy as np
from PIL import Image, ImageOps
import gradio as gr
import torch
from diffusers import DiffusionPipeline, EulerDiscreteScheduler

# ---------- CONFIG ----------
MODEL_ID = os.getenv("MODEL_ID", "hakurei/waifu-diffusion")  # change if needed
HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN", None)

device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32

# CPU-friendly defaults & limits
DEFAULT_WIDTH = 384
DEFAULT_HEIGHT = 384
DEFAULT_STEPS = 10
DEFAULT_GUIDANCE = 5.5
MAX_SEED = np.iinfo(np.int32).max

# ---------- Load pipeline (lazy) ----------
PIPE = None
def load_pipeline():
    global PIPE
    if PIPE is not None:
        return PIPE
    try:
        pipe = DiffusionPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=torch_dtype,
            use_auth_token=HF_TOKEN,
        )
        # Try to set EulerDiscreteScheduler if provided by model repo
        try:
            scheduler = EulerDiscreteScheduler.from_pretrained(MODEL_ID, subfolder="scheduler")
            pipe.scheduler = scheduler
        except Exception:
            pass

        pipe = pipe.to(device)

        # Optional: disable safety checker on CPU for speed (non-ideal but common)
        try:
            pipe.safety_checker = None
        except Exception:
            pass

        PIPE = pipe
        return PIPE
    except Exception as e:
        raise RuntimeError(f"Model load failed: {e}")

# ---------- Helpers ----------
DEFAULT_NEG = (
    "low quality, bad anatomy, blurry, extra limbs, malformed, deformed, "
    "watermark, text, signature, lowres"
)

def tidy_image(img: Image.Image, max_side=1024):
    img = img.convert("RGB")
    if max(img.size) > max_side:
        img = ImageOps.contain(img, (max_side, max_side))
    return img

# ---------- Inference ----------
def infer(
    prompt: str,
    negative_prompt: str,
    seed: int,
    randomize_seed: bool,
    width: int,
    height: int,
    guidance_scale: float,
    num_inference_steps: int,
):
    start = time.time()
    if not prompt or prompt.strip() == "":
        return None, "Enter a prompt."

    if randomize_seed or int(seed) == 0:
        seed = random.randint(0, MAX_SEED)
    else:
        seed = int(seed) % MAX_SEED

    try:
        pipe = load_pipeline()
    except Exception as e:
        return None, f"Model load error: {e}"

    # enforce CPU-friendly caps
    width = int(min(max(256, width), 512))
    height = int(min(max(256, height), 512))
    steps = int(min(max(4, num_inference_steps), 20))

    gen = torch.Generator(device=device).manual_seed(seed)

    try:
        out = pipe(
            prompt=prompt,
            negative_prompt=(negative_prompt or DEFAULT_NEG),
            width=width,
            height=height,
            guidance_scale=float(guidance_scale),
            num_inference_steps=steps,
            generator=gen,
        )
        image = tidy_image(out.images[0], max_side=1024)
        elapsed = time.time() - start
        return image, f"Done — Seed: {seed} • {int(elapsed)}s"
    except Exception:
        # lighter retry
        try:
            out = pipe(
                prompt=prompt,
                negative_prompt=(negative_prompt or DEFAULT_NEG),
                width=width,
                height=height,
                guidance_scale=max(3.0, float(guidance_scale) - 1.0),
                num_inference_steps=max(4, steps - 4),
                generator=gen,
            )
            image = tidy_image(out.images[0], max_side=1024)
            elapsed = time.time() - start
            return image, f"Recovered (retry) — Seed: {seed} • {int(elapsed)}s"
        except Exception as e2:
            return None, f"Generation failed: {e2}"

# ---------- UI (compatible with older/newer Gradio) ----------
css = """
/* Vibrant purple-pink gradient background */
body, .gradio-container {
  background: linear-gradient(135deg, #d946ef 0%, #a855f7 25%, #8b5cf6 50%, #7c3aed 75%, #6366f1 100%) !important;
  font-family: "Inter", system-ui, -apple-system, "Segoe UI", Roboto, "Helvetica Neue", Arial;
  color: #ffffff !important;
  min-height: 100vh;
}

/* Main container styling */
.contain, .gradio-container > div {
  background: transparent !important;
}

/* Header card - bright gradient with glow */
.header {
  padding: 20px 24px;
  border-radius: 16px;
  background: linear-gradient(135deg, #ec4899 0%, #f97316 50%, #facc15 100%);
  color: white;
  box-shadow: 0 8px 32px rgba(236, 72, 153, 0.4), 0 0 60px rgba(249, 115, 22, 0.3);
  margin-bottom: 20px;
}

/* Brand/title */
.brand {
  font-weight: 900;
  font-size: 28px;
  letter-spacing: 0.5px;
  color: #fff;
  text-shadow: 2px 2px 8px rgba(0,0,0,0.3);
}

/* Subtitle under brand */
.small {
  font-size: 14px;
  color: rgba(255,255,255,0.95);
  margin-top: 8px;
  font-weight: 500;
}

/* All blocks and containers - vibrant semi-transparent cards */
.gr-block, .gr-box, .gr-form, .gr-panel {
  background: rgba(255, 255, 255, 0.15) !important;
  backdrop-filter: blur(10px) !important;
  border-radius: 16px !important;
  border: 2px solid rgba(255, 255, 255, 0.25) !important;
  box-shadow: 0 8px 32px rgba(0, 0, 0, 0.1) !important;
  padding: 16px !important;
}

/* Input fields - bright with good contrast */
.gr-textbox, .gr-input, textarea, input {
  background: rgba(255, 255, 255, 0.95) !important;
  color: #1f2937 !important;
  border: 2px solid rgba(236, 72, 153, 0.3) !important;
  border-radius: 12px !important;
  padding: 12px !important;
  font-size: 15px !important;
  font-weight: 500 !important;
}

.gr-textbox::placeholder, textarea::placeholder, input::placeholder {
  color: rgba(31, 41, 55, 0.5) !important;
}

/* Labels - bright and visible */
label, .gr-label, .gr-box label {
  color: #ffffff !important;
  font-weight: 700 !important;
  font-size: 15px !important;
  text-shadow: 1px 1px 3px rgba(0,0,0,0.3) !important;
  margin-bottom: 8px !important;
}

/* Buttons - super vibrant gradient */
.gr-button, button {
  background: linear-gradient(135deg, #ff0080 0%, #ff8c00 50%, #ffd700 100%) !important;
  color: white !important;
  font-weight: 800 !important;
  border: none !important;
  box-shadow: 0 6px 24px rgba(255, 0, 128, 0.4), 0 0 40px rgba(255, 140, 0, 0.3) !important;
  border-radius: 12px !important;
  padding: 14px 24px !important;
  font-size: 16px !important;
  text-transform: uppercase;
  letter-spacing: 0.5px;
  transition: all 0.3s ease !important;
}

.gr-button:hover, button:hover {
  transform: translateY(-2px);
  box-shadow: 0 8px 32px rgba(255, 0, 128, 0.6), 0 0 60px rgba(255, 140, 0, 0.5) !important;
}

/* Sliders - bright colors */
.gr-slider input[type="range"] {
  background: rgba(255, 255, 255, 0.2) !important;
}

.gr-slider input[type="range"]::-webkit-slider-thumb {
  background: linear-gradient(135deg, #ff0080, #ff8c00) !important;
  border: 3px solid white !important;
  box-shadow: 0 2px 8px rgba(0,0,0,0.3) !important;
}

/* Accordion - vibrant */
.gr-accordion {
  background: rgba(255, 255, 255, 0.1) !important;
  border: 2px solid rgba(255, 255, 255, 0.2) !important;
  border-radius: 12px !important;
}

.gr-accordion summary {
  color: #ffffff !important;
  font-weight: 700 !important;
  background: rgba(236, 72, 153, 0.3) !important;
  padding: 12px !important;
  border-radius: 10px !important;
}

/* Image containers - bright white background */
.gr-image, .gr-gallery {
  background: rgba(255, 255, 255, 0.95) !important;
  border-radius: 12px !important;
  padding: 12px !important;
  border: 2px solid rgba(236, 72, 153, 0.3) !important;
}

/* Status textbox - bright and visible */
.gr-textbox[aria-label="Status"] {
  background: rgba(255, 255, 255, 0.9) !important;
  color: #1f2937 !important;
  border: 2px solid rgba(16, 185, 129, 0.5) !important;
  font-weight: 600 !important;
}

/* Examples - vibrant cards */
.gr-examples {
  background: rgba(255, 255, 255, 0.1) !important;
  border-radius: 12px !important;
  padding: 12px !important;
}

.gr-examples .gr-button {
  background: rgba(139, 92, 246, 0.8) !important;
  font-size: 13px !important;
  padding: 10px 16px !important;
}

/* Checkbox styling */
.gr-checkbox {
  color: #ffffff !important;
}

.gr-checkbox input[type="checkbox"] {
  border: 2px solid rgba(255, 255, 255, 0.5) !important;
  background: rgba(255, 255, 255, 0.2) !important;
}

/* Number inputs */
.gr-number input {
  background: rgba(255, 255, 255, 0.95) !important;
  color: #1f2937 !important;
}

/* Markdown text */
.gr-markdown, .markdown {
  color: #ffffff !important;
}

.gr-markdown strong {
  color: #fbbf24 !important;
  font-weight: 800 !important;
}

/* Mobile-friendly adjustments */
@media (max-width: 720px) {
  .header { text-align: center; }
  .brand { font-size: 24px; }
  .gr-button, button { font-size: 14px !important; padding: 12px 20px !important; }
}
"""


examples = [
    "anime girl standing on a cherry-blossom bridge at sunset, cinematic lighting, detailed eyes",
    "young samurai on a misty mountain path, dramatic clouds, anime style",
    "cozy studio apartment with anime character reading by window, warm lighting",
]

with gr.Blocks(css=css, title="MangaMorph — Anime Scene Generator") as demo:
    with gr.Row():
        with gr.Column(scale=2):
            gr.Markdown(
                "<div class='header'><div class='brand'>MangaMorph</div>"
                "<div class='small'>Text → Anime image • CPU-optimized • Try 384×384 & 10 steps for speed</div></div>"
            )
            prompt = gr.Textbox(lines=3, label="Describe your anime scene", placeholder="e.g. A cyberpunk anime girl on a rainy street...")
            with gr.Row():
                run_btn = gr.Button("Generate")
                download_btn = gr.Button("Download")
            with gr.Accordion("Advanced settings", open=False):
                negative = gr.Textbox(lines=2, label="Negative prompt (optional)", value=DEFAULT_NEG)
                with gr.Row():
                    seed = gr.Number(label="Seed (0 = random)", value=0)
                    randomize = gr.Checkbox(label="Randomize seed", value=True)
                with gr.Row():
                    width = gr.Slider(label="Width", minimum=256, maximum=512, step=64, value=DEFAULT_WIDTH)
                    height = gr.Slider(label="Height", minimum=256, maximum=512, step=64, value=DEFAULT_HEIGHT)
                with gr.Row():
                    guidance = gr.Slider(label="Guidance scale", minimum=1.0, maximum=12.0, step=0.1, value=DEFAULT_GUIDANCE)
                    steps = gr.Slider(label="Steps", minimum=4, maximum=20, step=1, value=DEFAULT_STEPS)
            gr.Examples(examples=examples, inputs=[prompt], label="Try examples")
            status = gr.Textbox(label="Status", value="Ready", interactive=False)

        with gr.Column(scale=1):
            gr.Markdown("**Preview**")
            result = gr.Image(label="Generated image")
            gallery = gr.Gallery(label="History (latest first)", columns=1)
            gr.Markdown("<div style='font-size:12px;color:#fff;margin-top:6px;font-weight:600;'>Tip: Use lower resolution & fewer steps for faster results on CPU</div>")

    def generate_and_update(prompt_text, negative_prompt_text, seed_val, randomize_val, w, h, g, s):
        img, msg = infer(prompt_text, negative_prompt_text, seed_val, randomize_val, w, h, g, s)
        history = [] if img is None else [img]
        return img, msg, history

    run_btn.click(
        fn=generate_and_update,
        inputs=[prompt, negative, seed, randomize, width, height, guidance, steps],
        outputs=[result, status, gallery],
        show_progress=True,
    )

    def download_current(image):
        # return image to trigger download
        return image

    download_btn.click(fn=download_current, inputs=[result], outputs=[result])

if __name__ == "__main__":
    demo.launch()