import gradio as gr
import openai
import requests
from PIL import Image
from io import BytesIO
import textwrap

def create_dalle_prompts(story_title, paragraph):
    return f"Illustration of '{story_title}' - {textwrap.shorten(paragraph, width=50)}"

def generate_image(api_key, prompt):
    openai.api_key = api_key
    response = openai.Image.create(
        prompt=prompt,
        n=1,
        size="1024x1024"
    )
    image_url = response['data'][0]['url']
    response = requests.get(image_url)
    image = Image.open(BytesIO(response.content))
    return image

def generate_story_and_images(name, story_title, story_type, api_key):
    openai.api_key = api_key
    if story_type == "for children":
        prompt_content = f"In a mystical kingdom, a brave adventurer named {name} discovers a magical object. Narrate their whimsical and heartwarming tale based on the title: '{story_title}'."
    else:
        prompt_content = f"In an odd city, {name} stumbles upon a series of bizarre events linked to the title '{story_title}'. Describe their hilariously absurd adventure."
    
    messages = [
        {"role": "system", "content": "You are a magical storyteller."},
        {"role": "user", "content": prompt_content}
    ]
    
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=messages
    )
    
    story = response.choices[0].message['content']
    
    # Split story into four paragraphs
    paragraphs = story.split('\n')[:4]
    
    # Create DALLE prompts for each paragraph
    dalle_prompts = [create_dalle_prompts(story_title, para) for para in paragraphs]
    
    images = [generate_image(api_key, prompt) for prompt in dalle_prompts]
    return story, dalle_prompts, *images

iface = gr.Interface(
    fn=generate_story_and_images,
    inputs=[
        gr.components.Textbox(label="Enter name (Example: Ella)"),
        gr.components.Textbox(label="Enter story title (Example: The Enchanted Locket)"),
        gr.components.Radio(["for children", "absurd"], label="Story Type"),
        gr.components.Textbox(label="OpenAI API Key", type="password")
    ],
    outputs=[
        "text",
        "text",
        gr.components.Image(label="Image 1", type="pil"),
        gr.components.Image(label="Image 2", type="pil"),
        gr.components.Image(label="Image 3", type="pil"),
        gr.components.Image(label="Image 4", type="pil")
    ],
    live=True
)

iface.launch()