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 remove_background(image):
    # Convert NumPy array to PIL Image
    #image = Image.fromarray(image_array)
    # 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)

    result = client.predict(images=handle_file(filename), api_name="/image")
    return Image.open(result[0])

def superimpose(image_with_text, overlay_image):
    # Open both images as RGBA to handle transparency
    #background = image_with_text #Image.open("/content/out2.png").convert("RGBA")
    overlay_image = overlay_image.convert("RGBA")
    # Paste overlay on the background
    image_with_text.paste(overlay_image, (0, 0), overlay_image)

    # Display or save the final image
    #background.show()
    # background.save("output_image.png")
    return image_with_text


def add_text_to_image(
    input_image,
    text,
    font_size,
    color,
    opacity,
    x_position,
    y_position
):
    """
    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))

    # Draw the main text
    draw.text(
        (actual_x, actual_y),
        text,
        font=font,
        fill=(*rgb_color, int(opacity))
    )

    # 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(title="Text on Image Editor",) as app:
        gr.Markdown("# Add Text to 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")
                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
            ],
            outputs=output_image
        )
        
        # Add example inputs
        gr.Examples(
            examples=[
                [
                    "sample.jpg",  # Replace with actual example image path
                    "SAMPLE TEXT",
                    400,
                    "White",
                    255,
                    50,
                    50
                ]
            ],
            inputs=[
                input_image,
                text_input,
                font_size,
                color_dropdown,
                opacity_slider,
                x_position,
                y_position
            ]
        )

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