import gradio as gr
import requests
from PIL import Image, ImageDraw, ImageFont
import textwrap

# Function to download image based on prompt
def download_image(prompt, index):
    url = f"https://pollinations.ai/p/{prompt}"
    response = requests.get(url)
    image_filename = f'generated_image_{index}.jpg'
    with open(image_filename, 'wb') as file:
        file.write(response.content)
    return image_filename

# Function to add text to image
def add_text_to_image(image_path, text):
    img = Image.open(image_path)
    draw = ImageDraw.Draw(img)
    font = ImageFont.load_default()  # You can load a custom font if you want
    # Wrap text for multiple lines
    lines = textwrap.fill(text, width=40)
    draw.text((10, 10), lines, font=font, fill="white")
    img_with_text = f'text_added_{image_path}'
    img.save(img_with_text)
    return img_with_text

# Function to visualize each line of the story
def visualize_story_lines(story):
    lines = story.split('\n')  # Split story into lines
    images_with_text = []
    
    for idx, line in enumerate(lines):
        prompt = line.replace(" ", "_")  # You can adjust prompt formatting
        img_file = download_image(prompt, idx)
        img_with_text = add_text_to_image(img_file, line)  # Add line as text to image
        images_with_text.append(img_with_text)
    
    return images_with_text  # Return list of images with text

# Gradio interface
def visualize_story_images(story):
    return visualize_story_lines(story)

iface = gr.Interface(fn=visualize_story_images, inputs="text", outputs="image")
iface.launch()