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

# Improved function for image generation with better text handling
def generate_image(text_description):
    # Create a blank image with a random background color
    img = Image.new('RGB', (512, 512), color=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
    draw = ImageDraw.Draw(img)
    
    # Load a default font (fallback to system font if arial.ttf is unavailable)
    try:
        font = ImageFont.truetype("arial.ttf", 30)
    except:
        font = ImageFont.load_default()
    
    # Wrap text to fit within the image
    wrapper = textwrap.TextWrapper(width=20)  # Adjust width for better wrapping
    wrapped_text = wrapper.fill(text=text_description)
    lines = wrapped_text.split('\n')
    
    # Draw each line of text
    y_text = 20
    for line in lines:
        line_width, line_height = draw.textsize(line, font=font)
        draw.text((20, y_text), line, font=font, fill=(255, 255, 255))
        y_text += line_height + 5  # Add spacing between lines
    
    # Optional: Add a simple graphical element to represent the description
    if "lake" in text_description.lower() or "mountains" in text_description.lower():
        draw.rectangle([200, 300, 300, 400], fill=(0, 191, 255))  # Blue rectangle for lake
        draw.polygon([(250, 250), (300, 200), (350, 250)], fill=(255, 255, 255))  # White triangle for mountain
    
    return img

# Gradio interface
with gr.Blocks(title="Text-to-Image Generator") as demo:
    gr.Markdown("# Text-to-Image Generator")
    gr.Markdown("Enter a description below and generate an image!")
    
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(label="Description", placeholder="Type your image description here...", value="A serene lake surrounded by snow-capped mountains under a vibrant sunset sky with shades of orange and purple.")
            generate_btn = gr.Button("Generate Image")
        with gr.Column():
            output_image = gr.Image(label="Generated Image")
    
    # Connect the button to the function
    generate_btn.click(fn=generate_image, inputs=text_input, outputs=output_image)

# Launch the app
demo.launch()