import gradio as gr
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import os
from PIL import Image
from gradio_client import Client, handle_file
import uuid

client = Client("ysharma/BiRefNet_for_text_writing")

def add_text_with_stroke(draw, text, x, y, font, text_color, stroke_width):
    """Helper function to draw text with stroke"""
    # Draw the stroke/outline
    for adj_x in range(-stroke_width, stroke_width + 1):
        for adj_y in range(-stroke_width, stroke_width + 1):
            draw.text((x + adj_x, y + adj_y), text, font=font, fill=text_color)

def remove_background(image):
    # Save the image to a specific location
    filename = f"image_{uuid.uuid4()}.png"  # Generates a universally unique identifier (UUID) for the filename
    image.save(filename)
    # Call gradio client for background removal
    result = client.predict(images=handle_file(filename), api_name="/image")
    return Image.open(result[0])

def superimpose(image_with_text, overlay_image):
    # Open image as RGBA to handle transparency
    overlay_image = overlay_image.convert("RGBA")
    # Paste overlay on the background
    image_with_text.paste(overlay_image, (0, 0), overlay_image)
    # Save the final image
    # image_with_text.save("output_image.png")
    return image_with_text

def add_text_to_image(
    input_image,
    text,
    font_size,
    color,
    opacity,
    x_position,
    y_position,
    thickness
):
    """
    Add text to an image with customizable properties
    """
    # Convert gradio image (numpy array) to PIL Image
    if input_image is None:
        return None
    
    image = Image.fromarray(input_image)
    # remove background 
    overlay_image = remove_background(image)

    # Create a transparent overlay for the text
    txt_overlay = Image.new('RGBA', image.size, (255, 255, 255, 0))
    draw = ImageDraw.Draw(txt_overlay)
    
    # Create a font with specified size
    try:
        font = ImageFont.truetype("DejaVuSans.ttf", int(font_size))
    except:
        # If DejaVu font is not found, try to use Arial or default
        try:
            font = ImageFont.truetype("arial.ttf", int(font_size))
        except:
            print("Using default font as system fonts not found")
            font = ImageFont.load_default()
    
    # Convert color name to RGB
    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),
        'Purple': (128, 0, 128)
    }
    rgb_color = color_map.get(color, (255, 255, 255))
    
    # Get text size for positioning
    text_bbox = draw.textbbox((0, 0), text, font=font)
    text_width = text_bbox[2] - text_bbox[0]
    text_height = text_bbox[3] - text_bbox[1]

    # Calculate actual x and y positions based on percentages
    actual_x = int((image.width - text_width) * (x_position / 100))
    actual_y = int((image.height - text_height) * (y_position / 100))

    # Create final color with opacity
    text_color = (*rgb_color, int(opacity))
    
    # Draw the text with stroke for thickness
    add_text_with_stroke(
        draw, 
        text, 
        actual_x, 
        actual_y, 
        font, 
        text_color, 
        int(thickness)
    )

    # Combine the original image with the text overlay
    if image.mode != 'RGBA':
        image = image.convert('RGBA')
    output_image = Image.alpha_composite(image, txt_overlay)
    
    # Convert back to RGB for display
    output_image = output_image.convert('RGB')

    # superimpose images
    output_image = superimpose(output_image, overlay_image)
    
    # Convert PIL image back to numpy array for Gradio
    return np.array(output_image)

# Create the Gradio interface
def create_interface():
    with gr.Blocks() as app:
        gr.Markdown("# Add Text Behind Image")
        gr.Markdown("Upload an image and customize text properties to add text overlay.")
        
        with gr.Row():
            with gr.Column():
                # Input components
                input_image = gr.Image(label="Upload Image", type="numpy")
                text_input = gr.Textbox(label="Enter Text", placeholder="Type your text here...")
                font_size = gr.Slider(minimum=10, maximum=800, value=400, step=10, 
                                    label="Font Size")
                thickness = gr.Slider(minimum=0, maximum=20, value=0, step=1,
                                    label="Text Thickness")
                color_dropdown = gr.Dropdown(
                    choices=["White", "Black", "Red", "Green", "Blue", "Yellow", "Purple"],
                    value="White",
                    label="Text Color"
                )
                opacity_slider = gr.Slider(minimum=0, maximum=255, value=255, step=1, 
                                         label="Opacity")
                x_position = gr.Slider(minimum=0, maximum=100, value=50, step=1, 
                                     label="X Position (%)")
                y_position = gr.Slider(minimum=0, maximum=100, value=50, step=1, 
                                     label="Y Position (%)")
                
            with gr.Column():
                # Output image
                output_image = gr.Image(label="Output Image")
        
        # Process button
        process_btn = gr.Button("Add Text to Image")
        
        # Connect the input components to the processing function
        process_btn.click(
            fn=add_text_to_image,
            inputs=[
                input_image,
                text_input,
                font_size,
                color_dropdown,
                opacity_slider,
                x_position,
                y_position,
                thickness
            ],
            outputs=output_image
        )
        
        # Add example inputs
        gr.Examples(
            examples=[
                [
                    "pink_convertible.webp",
                    "EPIC",
                    420,
                    "Purple",
                    150,
                    50,
                    21,
                    9
                ],
                [
                    "pear.jpg",
                    "PEAR",
                    350,
                    "Black",
                    100,
                    50,
                    2,
                    5
                ],
                [
                    "sample_text_image.jpeg",
                    "LIFE",
                    400,
                    "Black",
                    150,
                    50,
                    2,
                    8
                ],
            ],
            inputs=[
                input_image,
                text_input,
                font_size,
                color_dropdown,
                opacity_slider,
                x_position,
                y_position,
                thickness
            ],
            outputs=output_image,
            fn=add_text_to_image,
            cache_examples=True,
        )

    return app

# Launch the app
if __name__ == "__main__":
    # Try to install required font
    try:
        import subprocess
        subprocess.run(['apt-get', 'update'])
        subprocess.run(['apt-get', 'install', '-y', 'fonts-dejavu'])
        print("Font installed successfully")
    except:
        print("Could not install font automatically. Please install DejaVu font manually.")
    
    # Create and launch the interface
    app = create_interface()
    app.launch()