import spaces  # MUST come before torch / any CUDA-touching import

import random
import time

import torch
import gradio as gr
from inference import ShellDInference

MODEL_ID = "FlameF0X/ShellD"
TEXT_ENCODER = "sentence-transformers/all-MiniLM-L6-v2"

# Load at module scope, .to("cuda") eagerly. ShellDInference auto-detects cuda
# (the spaces hijack patches torch.cuda.is_available() in the main process).
# Pass the text encoder explicitly — the config.json's text_encoder_name is a
# local path ("./all-MiniLM-L6-v2") which doesn't exist in the Space.
pipe = ShellDInference(MODEL_ID, device="cpu", text_encoder_path=TEXT_ENCODER)
pipe.model.eval()
print("ShellD model loaded.")


def _estimate_duration(prompt, seed, randomize_seed, num_steps, cfg_scale, *args, **kwargs):
    """GPU-seconds reservation for ZeroGPU. Measured on ZeroGPU: ~0.01 s/step
    for this 67M-param DiT (2 CFG forward passes per step on 64 patches).
    Add ~5 s for cold-start weight streaming, floor at 10 s, ceiling at 60 s."""
    return min(60, max(10, int(num_steps * 0.015 + 5)))


@spaces.GPU(duration=_estimate_duration)
def generate(
    prompt: str,
    seed: int = 42,
    randomize_seed: bool = False,
    num_steps: int = 250,
    cfg_scale: float = 3.0,
    progress=gr.Progress(track_tqdm=False),
):
    """Generate a 256×256 image from a text prompt using ShellD.

    Args:
        prompt: Text description of the image to generate.
        seed: RNG seed for reproducibility (ignored if randomize_seed is True).
        randomize_seed: If True, pick a random seed and write it back.
        num_steps: Number of DDPM denoising steps (more = higher quality, slower).
        cfg_scale: Classifier-free guidance scale (higher = more prompt adherence).

    Returns:
        (image, seed) — a 256×256 PIL Image and the seed used.
    """
    if prompt is None or str(prompt).strip() == "":
        prompt = "a serene lake surrounded by mountains"
    prompt = str(prompt).strip()

    if randomize_seed:
        seed = random.randint(0, 2**31 - 1)
    seed = int(seed)

    t0 = time.perf_counter()
    img = pipe.generate(
        prompt=prompt,
        num_steps=int(num_steps),
        cfg_scale=float(cfg_scale),
        seed=seed,
    )
    elapsed = time.perf_counter() - t0
    print(f"ShellD generated '{prompt[:40]}' in {elapsed:.2f}s ({num_steps} steps, seed={seed})")
    return img, seed


CSS = """
#col-container { max-width: 900px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(css=CSS) as demo:
    gr.Markdown(
        """
        # 🐚 ShellD — Shell Diffusion Text-to-Image
        A lightweight 67M-parameter Diffusion Transformer (DiT) that generates 256×256 images from text prompts.
        [Model card](https://huggingface.co/FlameF0X/ShellD)
        """
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            prompt = gr.Textbox(
                label="Prompt",
                show_label=False,
                placeholder="Describe the image you want to generate…",
                container=False,
                scale=4,
            )
            run = gr.Button("Generate", variant="primary", scale=1)

        output = gr.Image(label="Generated image", height=320, show_label=True)

        with gr.Accordion("Advanced settings", open=False):
            num_steps = gr.Slider(
                label="Steps", minimum=10, maximum=1000, step=10, value=250
            )
            cfg_scale = gr.Slider(
                label="CFG scale", minimum=1.0, maximum=10.0, step=0.5, value=3.0
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            seed = gr.Number(label="Seed", value=0, precision=0)

        gr.Examples(
            examples=[
                ["a serene lake surrounded by mountains"],
                ["a futuristic city at night with neon lights"],
                ["a cute cat sitting on a windowsill"],
                ["a tropical beach with palm trees at sunset"],
            ],
            inputs=[prompt],
            outputs=[output, seed],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        generate,
        inputs=[prompt, seed, randomize_seed, num_steps, cfg_scale],
        outputs=[output, seed],
        api_name="generate",
    )
    prompt.submit(
        generate,
        inputs=[prompt, seed, randomize_seed, num_steps, cfg_scale],
        outputs=[output, seed],
        api_name="generate_submit",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus())