import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFont

# Define a dictionary to map color names to their RGB values
COLOR_MAP = {
    "white": (255, 255, 255),
    "black": (0, 0, 0),
    "red": (255, 0, 0),
    "green": (0, 255, 0),
    "blue": (0, 0, 255),
    "yellow": (255, 255, 0),
    "cyan": (0, 255, 255),
    "magenta": (255, 0, 255),
}

# Define the function to add text to the image
def add_text_to_image(image, text, text_color, text_size, width, height, text_align, font):
    try:
        # Convert the image to a PIL Image object if it's not already
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image.astype('uint8'), 'RGB')
        
        # Resize the image to the specified width and height while maintaining aspect ratio
        original_width, original_height = image.size
        aspect_ratio = original_width / original_height

        # Calculate new dimensions while maintaining aspect ratio
        if width / height > aspect_ratio:
            new_width = int(height * aspect_ratio)
            new_height = height
        else:
            new_width = width
            new_height = int(width / aspect_ratio)

        # Resize the image
        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)

        # Create a new blank image with the specified width and height
        new_image = Image.new("RGB", (width, height), (255, 255, 255))  # White background
        # Paste the resized image onto the new image, centered
        new_image.paste(image, ((width - new_width) // 2, (height - new_height) // 2))

        # Create a drawing context
        draw = ImageDraw.Draw(new_image)

        # Load the font (you can specify a path to a .ttf file if needed)
        try:
            font = ImageFont.truetype(font, text_size)
        except IOError:
            # Fallback to a default font if the specified font is not available
            font = ImageFont.load_default()

        # Get the RGB color from the COLOR_MAP
        color_rgb = COLOR_MAP.get(text_color, (255, 255, 255))  # Default to white if color not found

        # Calculate text position based on alignment
        text_width, text_height = draw.textsize(text, font=font)
        if text_align == "left":
            text_position_x = 10  # 10 pixels from the left edge
        elif text_align == "right":
            text_position_x = width - text_width - 10  # 10 pixels from the right edge
        elif text_align == "center":
            text_position_x = (width - text_width) // 2  # Centered horizontally

        text_position_y = (height - text_height) // 2  # Centered vertically by default

        # Add the text to the image
        draw.text((text_position_x, text_position_y), text, fill=color_rgb, font=font)

        # Convert the image back to a numpy array for Gradio compatibility
        return np.array(new_image)

    except Exception as e:
        # Handle any errors and display a message
        print(f"Error: {e}")
        return None

# Gradio interface
iface = gr.Interface(
    fn=add_text_to_image,
    inputs=[
        gr.Image(label="Input Image"),
        gr.Textbox(label="Text"),
        gr.Dropdown(label="Text Color", choices=list(COLOR_MAP.keys()), value="white"),
        gr.Slider(minimum=10, maximum=100, value=30, label="Text Size"),
        gr.Slider(minimum=100, maximum=2048, value=1024, label="Width"),
        gr.Slider(minimum=100, maximum=2048, value=1024, label="Height"),
        gr.Dropdown(label="Text Alignment", choices=["left", "right", "center"], value="center"),
        gr.Textbox(label="Font (e.g., 'arial.ttf')", value="arial.ttf")
    ],
    outputs=gr.Image(label="Output Image"),
    title="Add Text to Image",
    description="Upload an image, resize it, and add customizable text with alignment options."
)

# Launch the Gradio app
iface.launch()