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

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

def add_text_to_image(image_path, text):
    try:
        img = Image.open(image_path)
        draw = ImageDraw.Draw(img)
        
        # Use a default font if custom font is not available
        try:
            font = ImageFont.truetype("arial.ttf", 30)
        except IOError:
            font = ImageFont.load_default().font_variant(size=30)
        
        # Calculate text size and position
        img_width, img_height = img.size
        max_text_width = int(img_width * 0.9)  # 90% of image width
        lines = textwrap.wrap(text, width=40)
        line_height = font.getsize('hg')[1] + 5  # Add some padding
        text_height = line_height * len(lines)
        
        # Create semi-transparent background
        overlay = Image.new('RGBA', img.size, (0, 0, 0, 0))
        overlay_draw = ImageDraw.Draw(overlay)
        overlay_draw.rectangle([(0, img_height - text_height - 20), (img_width, img_height)], 
                               fill=(0, 0, 0, 180))
        
        # Paste the overlay onto the original image
        img = img.convert('RGBA')
        img = Image.alpha_composite(img, overlay)
        draw = ImageDraw.Draw(img)
        
        # Add text
        y_text = img_height - text_height - 10
        for line in lines:
            line_width, _ = draw.textsize(line, font=font)
            x_text = (img_width - line_width) / 2
            draw.text((x_text, y_text), line, font=font, fill=(255, 255, 255, 255))
            y_text += line_height
        
        img_with_text = f'text_added_{os.path.basename(image_path)}'
        img = img.convert('RGB')
        img.save(img_with_text)
        return img_with_text
    except Exception as e:
        print(f"Error adding text to image: {e}")
        return None

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)]

def visualize_story_images(story):
    if not story.strip():
        return []
    return visualize_story_lines(story)

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)