# Required Imports
import streamlit as st
import google.generativeai as genai
import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
import io

# -----------------------------
# 🔐 CONFIGURE API KEYS
# -----------------------------
# Google Generative AI API Key
GOOGLE_API_KEY = "AIzaSyBn9Ehq5oIqkEov_fmAMT258X6imfNXfvg" 
genai.configure(api_key=GOOGLE_API_KEY)
model = genai.GenerativeModel("models/gemini-2.0-flash")

# -----------------------------
# 🎨 LOAD STABLE DIFFUSION MODEL
# -----------------------------
model_id = "prompthero/openjourney-v4"
device = "cuda" if torch.cuda.is_available() else "cpu"

try:
    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
    pipe.to(device)
    print("✅ Stable Diffusion Model Loaded Successfully!")
except Exception as e:
    print(f"❌ Error loading Stable Diffusion model: {e}")
    exit()

# -----------------------------
# 🧠 GENERATE STORY PANELS
# -----------------------------
def get_summary_or_concept(prompt, num_panels=6):
    content_prompt = (
        f"Summarize the key concepts or moments of '{prompt}' into {num_panels} brief points, "
        "with 1 concise sentence for each panel, equally distributing the information."
    )

    try:
        response = model.generate_content(content_prompt)
        if response and response.text:
            text_response = response.text.strip()
            panels = [p.strip() for p in text_response.split("\n") if p.strip()]
            if len(panels) >= num_panels:
                return panels[:num_panels]
            else:
                return panels + [panels[-1]] * (num_panels - len(panels))
    except Exception as e:
        print(f"❌ Error generating content: {e}")
        return ["Error: Unable to generate story."]

# -----------------------------
# 🎨 GENERATE COMIC IMAGES
# -----------------------------
def generate_comic_images(story_prompts, num_steps=20):
    images = []
    for i, panel in enumerate(story_prompts):
        st.write(f"🖼️ Generating image for Panel {i+1}/{len(story_prompts)}: {panel.strip()[:50]}...")

        image_prompt = (
            f"Highly detailed, cinematic scene depicting: {panel.strip()}, "
            "vivid colors, realistic style, no text, no words"
        )

        try:
            image = pipe(
                image_prompt,
                num_inference_steps=num_steps,
                guidance_scale=7.0,
                negative_prompt="text, words, letters, nsfw, explicit, unsafe"
            ).images[0]
            images.append(image)
        except Exception as e:
            st.error(f"❌ Error generating image: {e}")
            return None

    return images

# -----------------------------
# 🖼️ CREATE COMIC STRIP
# -----------------------------
def create_comic_strip(images, panel_width=300, spacing=10):
    widths, heights = zip(*(i.size for i in images))

    target_height = 300
    resized_images = [
        img.resize((int(w * target_height / h), target_height)) for img, (w, h) in zip(images, zip(widths, heights))
    ]

    total_width = sum(img.size[0] for img in resized_images) + (spacing * (len(images) - 1))
    comic_strip = Image.new("RGB", (total_width, target_height), "white")

    x_offset = 0
    for img in resized_images:
        comic_strip.paste(img, (x_offset, 0))
        x_offset += img.size[0] + spacing

    return comic_strip

# -----------------------------
#