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

# 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()  # Load default font
    # 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(" ", "_")  # Adjust prompt formatting
        img_file = download_image(prompt, idx)
        img_with_text = add_text_to_image(img_file, line)  # Add line text to image
        images_with_text.append(img_with_text)
    
    return images_with_text  # Return list of images with text

# Gradio interface with state management
def visualize_story_images(story, img_index):
    images = visualize_story_lines(story)
    if 0 <= img_index < len(images):
        return images[img_index]  # Return the image based on index
    return None

# Create a function for navigation
def navigate_story(story, img_index):
    images = visualize_story_lines(story)
    if img_index < len(images):
        return images[img_index]
    return images[-1]  # Return the last image if index exceeds

# Gradio interface setup
story_input = gr.inputs.Textbox(lines=10, placeholder="Enter your story here...")
image_output = gr.outputs.Image(type="pil")

iface = gr.Interface(
    fn=visualize_story_images,
    inputs=[story_input, gr.inputs.Slider(minimum=0, maximum=9, step=1, label="Image Index")],
    outputs=image_output,
    live=True
)

iface.launch(share=True)