import gradio as gr
from langchain.llms import OpenAI as LangChainOpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from openai import OpenAI
import os
import re

# Initialize OpenAI client
client = OpenAI()

# Create a prompt template
template = """
You are a creative story writer. Given a topic, write a short story of about 150 words. 
Divide the story into 3 paragraphs. Each paragraph should be a distinct part of the story.

Topic: {topic}

Story:
"""

prompt = PromptTemplate(template=template, input_variables=["topic"])

# Create an LLMChain
llm = LangChainOpenAI(temperature=0.7)
story_chain = LLMChain(llm=llm, prompt=prompt)

def generate_image(prompt):
    try:
        response = client.images.generate(
            model="dall-e-3",
            prompt=prompt,
            size="1024x1024",
            quality="standard",
            n=1,
        )
        return response.data[0].url
    except Exception as e:
        print(f"Error generating image: {e}")
        return None

def generate_story_with_images(topic):
    # Generate the story
    story = story_chain.run(topic)
    
    # Split the story into paragraphs
    paragraphs = re.split(r'\n\n', story.strip())
    
    # Ensure we have exactly 3 paragraphs
    paragraphs = paragraphs[:3]
    while len(paragraphs) < 3:
        paragraphs.append("...")

    # Generate images for each paragraph
    image_urls = []
    for paragraph in paragraphs:
        image_url = generate_image(paragraph)
        image_urls.append(image_url)
    
    return paragraphs, image_urls

# Create the Gradio interface
def gradio_interface(topic):
    try:
        paragraphs, image_urls = generate_story_with_images(topic)
        return [
            paragraphs[0], image_urls[0],
            paragraphs[1], image_urls[1],
            paragraphs[2], image_urls[2]
        ]
    except Exception as e:
        print(f"Error in gradio_interface: {e}")
        return ["An error occurred while generating the story and images."] + [""] * 5

iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Textbox(lines=2, placeholder="Enter a topic for the story..."),
    outputs=[
        gr.Textbox(label="Paragraph 1"),
        gr.Image(label="Image 1", type="filepath"),
        gr.Textbox(label="Paragraph 2"),
        gr.Image(label="Image 2", type="filepath"),
        gr.Textbox(label="Paragraph 3"),
        gr.Image(label="Image 3", type="filepath"),
    ],
    title="Story Generator with Images",
    description="Enter a topic, and the AI will generate a short story with an image for each paragraph."
)

# Launch the Gradio app
iface.launch()