import gradio as gr
import requests
from PIL import Image, ImageDraw, ImageFont
import textwrap
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

# Function to download image based on prompt
def download_image(prompt, index):
    try:
        url = f"https://pollinations.ai/p/{prompt}"
        response = requests.get(url, timeout=10)
        response.raise_for_status()
        image_filename = f'generated_image_{index}.jpg'
        with open(image_filename, 'wb') as file:
            file.write(response.content)
        return image_filename
    except requests.RequestException as e:
        print(f"Error downloading image {index}: {e}")
        return None

# Function to add text to image
def add_text_to_image(image_path, text):
    try:
        img = Image.open(image_path)
        draw = ImageDraw.Draw(img)
        # Use a default font if arial.ttf is not available
        try:
            font = ImageFont.truetype("arial.ttf", 20)
        except IOError:
            font = ImageFont.load_default()
        lines = textwrap.fill(text, width=40)
        draw.text((10, 10), lines, font=font, fill="white", stroke_width=2, stroke_fill="black")
        img_with_text = f'text_added_{os.path.basename(image_path)}'
        img.save(img_with_text)
        return img_with_text
    except Exception as e:
        print(f"Error adding text to image: {e}")
        return None

# Function to visualize each line of the story
def visualize_story_lines(story, progress=gr.Progress()):
    lines = [line.strip() for line in story.split('\n') if line.strip()]
    images_with_text = []
    
    def process_line(idx, line):
        prompt = line.replace(" ", "_")
        img_file = download_image(prompt, idx)
        if img_file:
            return add_text_to_image(img_file, line)
        return None

    with ThreadPoolExecutor(max_workers=5) as executor:
        future_to_idx = {executor.submit(process_line, idx, line): idx for idx, line in enumerate(lines)}
        for future in progress.tqdm(as_completed(future_to_idx), total=len(lines), desc="Processing images"):
            idx = future_to_idx[future]
            result = future.result()
            if result:
                images_with_text.append((idx, result))
    
    return [img for _, img in sorted(images_with_text)]

# Gradio interface with state management
def visualize_story_images(story):
    if not story.strip():
        return []
    return visualize_story_lines(story)

# Create a Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# Story Visualizer")
    
    with gr.Row():
        story_input = gr.Textbox(lines=10, placeholder="Enter your story here...", label="Story Input")
        visualize_button = gr.Button("Visualize Story")
    
    gallery_output = gr.Gallery(label="Visualized Story", show_label=False, elem_id="gallery", columns=2, rows=2, height="auto")
    
    visualize_button.click(
        fn=visualize_story_images,
        inputs=[story_input],
        outputs=[gallery_output]
    )

iface.launch(share=True)