import streamlit as st
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from PIL import Image
import io

# ── Page Config ───────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="Text to Image Generator",
    page_icon="🎨",
    layout="wide",
)

# ── Custom CSS ────────────────────────────────────────────────────────────────
st.markdown("""
<style>
    .main { background-color: #0f0f1a; }
    .stApp { background: linear-gradient(135deg, #0f0f1a 0%, #1a1a2e 100%); }
    h1 { color: #c084fc !important; font-size: 2.5rem !important; }
    h3 { color: #a78bfa !important; }
    .stButton > button {
        background: linear-gradient(90deg, #7c3aed, #a855f7);
        color: white;
        border: none;
        border-radius: 10px;
        padding: 0.6rem 2rem;
        font-size: 1rem;
        font-weight: 600;
        width: 100%;
        transition: opacity 0.2s;
    }
    .stButton > button:hover { opacity: 0.85; }
    .stTextArea textarea, .stTextInput input {
        background-color: #1e1e3a !important;
        color: #e2e8f0 !important;
        border: 1px solid #4c1d95 !important;
        border-radius: 8px !important;
    }
    .stSlider > div > div { color: #c084fc; }
    label { color: #c4b5fd !important; font-weight: 500; }
    .stImage img {
        border-radius: 12px;
        border: 1px solid #4c1d95;
        box-shadow: 0 0 20px rgba(168, 85, 247, 0.2);
    }
    .info-box {
        background: #1e1e3a;
        border-left: 4px solid #7c3aed;
        border-radius: 8px;
        padding: 12px 16px;
        margin-bottom: 16px;
        color: #c4b5fd;
        font-size: 0.9rem;
    }
    .stExpander { border: 1px solid #4c1d95 !important; border-radius: 8px !important; }
    .stDownloadButton > button {
        background: #1e1e3a !important;
        color: #c084fc !important;
        border: 1px solid #7c3aed !important;
        border-radius: 8px;
        width: 100%;
    }
</style>
""", unsafe_allow_html=True)


# ── Model Loader (cached) ─────────────────────────────────────────────────────
@st.cache_resource(show_spinner=False)
def load_pipeline():
    model_id = "runwayml/stable-diffusion-v1-5"
    pipe = StableDiffusionPipeline.from_pretrained(
        model_id,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        safety_checker=None,
        requires_safety_checker=False,
    )
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    pipe = pipe.to(device)
    if torch.cuda.is_available():
        pipe.enable_attention_slicing()
    return pipe


# ── Helper: PIL → bytes ───────────────────────────────────────────────────────
def image_to_bytes(img: Image.Image) -> bytes:
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


# ── UI ────────────────────────────────────────────────────────────────────────
st.title("🎨 Text to Image Generator")
st.markdown("**Powered by Stable Diffusion v1.5** — Transform your words into stunning visuals.")

col_left, col_right = st.columns([1, 1], gap="large")

with col_left:
    st.markdown("### ✏️ Describe your image")

    prompt = st.text_area(
        "Prompt",
        placeholder="a futuristic city at sunset, cyberpunk style, ultra detailed, 4k...",
        height=110,
        label_visibility="collapsed",
    )

    negative_prompt = st.text_area(
        "Negative Prompt",
        value="blurry, ugly, distorted, low quality, watermark, text, signature",
        height=75,
        help="Describe what you DON'T want in the image.",
    )

    with st.expander("⚙️ Advanced Settings"):
        col_w, col_h = st.columns(2)
        with col_w:
            width = st.select_slider("Width", options=[256, 320, 384, 448, 512, 576, 640, 704, 768], value=512)
        with col_h:
            height = st.select_slider("Height", options=[256, 320, 384, 448, 512, 576, 640, 704, 768], value=512)

        steps = st.slider("Inference Steps", min_value=10, max_value=50, value=25, step=1,
                          help="More steps = better quality but slower.")
        guidance = st.slider("Guidance Scale", min_value=1.0, max_value=15.0, value=7.5, step=0.5,
                             help="Higher = more literal to your prompt.")
        num_images = st.slider("Number of Images", min_value=1, max_value=4, value=1, step=1)
        seed = st.number_input("Seed (-1 = random)", value=-1, step=1,
                               help="Same seed + same prompt = same image every time.")

    st.markdown("### 💡 Example Prompts")
    examples = [
        "portrait of a samurai warrior, cinematic lighting, 4k",
        "dragon flying over snow-capped mountains, fantasy art",
        "astronaut riding a horse on Mars, photorealistic",
        "cute cartoon cat in sunglasses, vibrant illustration",
    ]
    for ex in examples:
        if st.button(f"📝 {ex[:45]}...", key=ex):
            st.session_state["example_prompt"] = ex

    if "example_prompt" in st.session_state:
        st.info(f"Copied: *{st.session_state['example_prompt']}* — paste it in the prompt box above.")

    generate_btn = st.button("✨ Generate Image")


with col_right:
    st.markdown("### 🖼️ Output")

    if generate_btn:
        if not prompt.strip():
            st.error("Please enter a prompt first!")
        else:
            with st.spinner("🔮 Loading model (first run takes ~2 min)..."):
                pipe = load_pipeline()

            with st.spinner(f"🎨 Generating {num_images} image(s)... (~{steps * 2}s)"):
                try:
                    generator = None
                    if seed != -1:
                        device = "cuda" if torch.cuda.is_available() else "cpu"
                        generator = torch.Generator(device).manual_seed(int(seed))

                    result = pipe(
                        prompt=prompt,
                        negative_prompt=negative_prompt if negative_prompt.strip() else None,
                        num_inference_steps=steps,
                        guidance_scale=guidance,
                        width=width,
                        height=height,
                        generator=generator,
                        num_images_per_prompt=num_images,
                    )

                    images = result.images
                    st.session_state["images"] = images
                    st.session_state["last_prompt"] = prompt
                    st.success(f"✅ Generated {len(images)} image(s)!")

                except Exception as e:
                    st.error(f"Generation failed: {e}")

    if "images" in st.session_state:
        images = st.session_state["images"]
        cols = st.columns(2 if len(images) > 1 else 1)
        for i, img in enumerate(images):
            with cols[i % len(cols)]:
                st.image(img, use_column_width=True)
                st.download_button(
                    label=f"⬇️ Download image {i + 1}",
                    data=image_to_bytes(img),
                    file_name=f"generated_{i + 1}.png",
                    mime="image/png",
                    key=f"dl_{i}",
                )

    else:
        st.markdown("""
        <div class="info-box">
        👈 Fill in your prompt on the left and hit <strong>Generate Image</strong> to get started.<br><br>
        🚀 <strong>Tips:</strong><br>
        • Add style words: <em>cinematic, 4k, oil painting, photorealistic</em><br>
        • Use negative prompts to remove unwanted elements<br>
        • Guidance scale 7–9 gives the best balance
        </div>
        """, unsafe_allow_html=True)

# ── Footer ────────────────────────────────────────────────────────────────────
st.divider()
st.markdown(
    "<p style='text-align:center; color:#6b7280; font-size:0.85rem;'>"
    "Built with Streamlit · Stable Diffusion v1.5 · Deployed on Hugging Face Spaces"
    "</p>",
    unsafe_allow_html=True,
)