import zipfile
from io import BytesIO

import gradio as gr
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
import torch
from PIL import Image, ImageDraw, ImageFont

# =====================================================
# Load Stable Diffusion with Euler Scheduler
# =====================================================
model_id = "stabilityai/stable-diffusion-2-1-base"
scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(
    model_id, scheduler=scheduler, torch_dtype=torch.float16
).to("cuda")  # Change to "cpu" if GPU unavailable

# =====================================================
# Comic Styles & Descriptions
# =====================================================
comic_styles = {
    "Superhero comics": "Bold, dynamic figures and vibrant colors. Emphasizes action, drama, and idealism.",
    "Alternative comics": "Independent comics outside the mainstream genre, from realistic to abstract.",
    "Noir": "High contrast light and shadow, deep blacks, gritty and moody atmosphere.",
    "Cartooning/Toon style": "Simplified, exaggerated, caricatured characters, e.g., Betty Boop.",
    "Retro comics/Silver Age": "1950s–1970s style, simpler energetic designs and vibrant colors.",
    "Webcomics": "Digital medium ranging from simple hand-drawn to polished sequential art.",
    "Shōnen": "For young male audience. Bold lines, dynamic action. Examples: Dragon Ball.",
    "Shōjo": "For young female readers. Fluid lines, detailed eyes, decorative aesthetic.",
    "Seinen": "For adult male readers. More realistic/gritty art, mature themes.",
    "Chibi": "Super-deformed cute characters with tiny bodies and oversized heads.",
    "Moe": "Cute, innocent-looking characters, appealing to fans.",
    "Kodomo": "For young children. Simple, rounded shapes, bright artwork.",
    "Line art": "Strong, distinct lines without extensive shading or coloring.",
    "Slice of life": "Depicts everyday events, cartoonish or realistic, relatable.",
    "Celtic art style": "Decorative, stylized, with intricate knotwork and patterns.",
    "Ink wash": "Diluted ink for soft, textured, or painted feel.",
    "Pointillism": "Small dots of color form images, creating textures.",
    "Watercolor painting style": "Soft, painterly look with translucent gradients.",
    "Photorealism": "Attempts photographic accuracy.",
    "Grotesque": "Distorted/exaggerated for weird, unsettling, or macabre effect."
}

# =====================================================
# Comic Generation Functions
# =====================================================
def add_speech_bubbles(img: Image.Image, dialogues):
    """Add simple speech bubbles at top-left positions."""
    draw = ImageDraw.Draw(img)
    try:
        font = ImageFont.truetype("ComicNeue-Bold.ttf", 20)
    except:
        font = ImageFont.load_default()
    x, y = 20, 20
    for text in dialogues:
        bubble_w, bubble_h = 200, 60
        rect = [x, y, x + bubble_w, y + bubble_h]
        draw.rectangle(rect, fill="white", outline="black", width=3)
        draw.text((x + 10, y + 10), text, font=font, fill="black")
        y += bubble_h + 10
    return img

def generate_panel(prompt):
    """Generate a single panel using Stable Diffusion."""
    return pipe(prompt).images[0]

def generate_comic(title, story, pages, panels, style, dialogues, explicit):
    """Generate multi-panel comics and return images + ZIP."""
    dialogues_list = [d.strip() for d in dialogues.splitlines() if d.strip()]
    generated_images = []

    for _ in range(int(pages)):
        for _ in range(int(panels)):
            prompt = f"{style} comic panel ({comic_styles[style]}), theme: {story}"
            if explicit:
                prompt += " with mature themes, blood and offensive language"
            else:
                prompt += " safe for all ages"

            img = generate_panel(prompt)

            if dialogues_list:
                img = add_speech_bubbles(img, dialogues_list)

            generated_images.append(img)

    # Save images to ZIP
    buffer = BytesIO()
    with zipfile.ZipFile(buffer, "w") as zf:
        for i, img in enumerate(generated_images):
            img_bytes = BytesIO()
            img.save(img_bytes, format="PNG")
            zf.writestr(f"comic_panel_{i+1}.png", img_bytes.getvalue())
    buffer.seek(0)

    return generated_images, (title or "comic") + ".zip", buffer

# =====================================================
# Gradio UI
# =====================================================
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("<h1 style='text-align:center'>🦸‍♂️ AI Comic Crafter</h1>"
                "<p style='text-align:center'>Generate multi-page, multi-panel comics with AI! 🎨</p>")

    with gr.Row():
        with gr.Column(scale=1):
            title = gr.Textbox(label="Comic Title", placeholder="Enter your comic book title...")
            story = gr.Textbox(label="Story Setup", placeholder="Describe characters, setting, and plot...")
            pages = gr.Slider(1, 5, value=2, step=1, label="Number of Pages")
            panels = gr.Slider(1, 6, value=3, step=1, label="Panels Per Page")
            
            style_dropdown = gr.Dropdown(
                choices=list(comic_styles.keys()),
                label="Art Style / Theme",
                value="Superhero comics"
            )
            style_description = gr.Textbox(
                value=comic_styles["Superhero comics"],
                label="Style Description",
                interactive=False,
                lines=5
            )
            style_dropdown.change(
                fn=lambda x: comic_styles[x],
                inputs=style_dropdown,
                outputs=style_description
            )

            dialogues = gr.Textbox(label="Custom Dialogues (Optional)", placeholder="One dialogue per line...", lines=5)
            explicit = gr.Checkbox(label="Allow Explicit Content (blood/gore, offensive language)", value=False)
            generate_btn = gr.Button("⚡ Generate Comic")

        with gr.Column(scale=1):
            output_gallery = gr.Gallery(label="Your Comic Preview", columns=2, height="auto")
            download_file = gr.File(label="Download Complete Comic (.zip)", type="binary")

    generate_btn.click(
        fn=generate_comic,
        inputs=[title, story, pages, panels, style_dropdown, dialogues, explicit],
        outputs=[output_gallery, download_file, download_file]
    )

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