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

# Constants for image styles
IMAGE_STYLES = {
    "Realistic": "",
    "Cartoon": "cartoon style",
    "Watercolor": "watercolor painting style",
    "Sketch": "pencil sketch style",
    "Anime": "anime style"
}

def download_image(prompt, index, style):
    try:
        style_prompt = f"{prompt}, {IMAGE_STYLES[style]}" if style != "Realistic" else prompt
        url = f"https://pollinations.ai/p/{style_prompt.replace(' ', '_')}"
        response = requests.get(url, timeout=15)
        response.raise_for_status()
        image_filename = f'generated_image_{index}_{int(time.time())}.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, font_size, text_color):
    try:
        img = Image.open(image_path)
        draw = ImageDraw.Draw(img)
        
        try:
            font = ImageFont.truetype("arial.ttf", font_size)
        except IOError:
            font = ImageFont.load_default().font_variant(size=font_size)
        
        img_width, img_height = img.size
        max_text_width = int(img_width * 0.9)
        lines = textwrap.wrap(text, width=40)
        
        line_height = font.getbbox("hg")[3] + 5
        text_height = line_height * len(lines)
        
        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))
        
        img = img.convert('RGBA')
        img = Image.alpha_composite(img, overlay)
        draw = ImageDraw.Draw(img)
        
        y_text = img_height - text_height - 10
        for line in lines:
            bbox = font.getbbox(line)
            line_width = bbox[2] - bbox[0]
            x_text = (img_width - line_width) / 2
            draw.text((x_text, y_text), line, font=font, fill=text_color)
            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 process_story_lines(story, style, font_size, text_color, 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, style)
        if img_file:
            return add_text_to_image(img_file, line, font_size, text_color)
        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, style, font_size, text_color):
    if not story.strip():
        return []
    return process_story_lines(story, style, font_size, text_color)

def preview_image(story, style, font_size, text_color):
    if not story.strip():
        return None
    lines = [line.strip() for line in story.split('\n') if line.strip()]
    if not lines:
        return None
    preview_line = random.choice(lines)
    img_file = download_image(preview_line, 0, style)
    if img_file:
        return add_text_to_image(img_file, preview_line, font_size, text_color)
    return None

with gr.Blocks(css="#gallery {height: 600px}") as iface:
    gr.Markdown("# 📚✨ Story Visualizer")
    
    with gr.Row():
        with gr.Column(scale=2):
            story_input = gr.Textbox(lines=10, placeholder="Enter your story here...", label="Story Input")
        with gr.Column(scale=1):
            style_dropdown = gr.Dropdown(choices=list(IMAGE_STYLES.keys()), value="Realistic", label="Image Style")
            font_size_slider = gr.Slider(minimum=20, maximum=60, value=30, step=2, label="Font Size")
            text_color_picker = gr.ColorPicker(value="#FFFFFF", label="Text Color")
            preview_button = gr.Button("Preview", variant="secondary")
            visualize_button = gr.Button("Visualize Story", variant="primary")
    
    preview_output = gr.Image(label="Preview", height=300)
    gallery_output = gr.Gallery(label="Visualized Story", show_label=False, elem_id="gallery", columns=2, rows=2, height="auto")
    
    preview_button.click(
        fn=preview_image,
        inputs=[story_input, style_dropdown, font_size_slider, text_color_picker],
        outputs=[preview_output]
    )
    
    visualize_button.click(
        fn=visualize_story_images,
        inputs=[story_input, style_dropdown, font_size_slider, text_color_picker],
        outputs=[gallery_output]
    )

    gr.Markdown("""
    ## How to use:
    1. Enter your story in the text box, with each line representing a scene.
    2. Choose an image style from the dropdown menu.
    3. Adjust the font size and text color as desired.
    4. Click "Preview" to see how a random scene from your story might look.
    5. When you're happy with the settings, click "Visualize Story" to generate images for all scenes.
    
    Enjoy bringing your story to life! 🎨📖
    """)

iface.launch(share=True)