import gradio as gr
import numpy as np
import PIL.Image as Image
from rembg import remove
from PIL import ImageDraw, ImageFont
import textwrap

# Function to remove background and place text behind person
def text_behind_image(input_image, text, text_color, font_size, text_opacity):
    if input_image is None or text.strip() == "":
        return input_image
    
    # Convert color hex to RGB
    try:
        r = int(text_color[1:3], 16)
        g = int(text_color[3:5], 16)
        b = int(text_color[5:7], 16)
        text_color_rgb = (r, g, b)
    except:
        text_color_rgb = (0, 0, 0)  # Default to black
    
    try:
        # Open the image
        img = Image.fromarray(input_image)
        
        # Get dimensions
        width, height = img.size
        
        # Step 1: Extract the person from the image
        try:
            # Try to use the human segmentation model first
            person_only = remove(img, model_name="u2net_human_seg")
        except:
            # Fallback to default model if human_seg not available
            person_only = remove(img)
            
        # Step 2: Create a mask from the person cutout
        person_mask = Image.new('RGBA', (width, height), (0, 0, 0, 0))
        person_mask.paste(person_only, (0, 0), person_only)
        
        # Step 3: Extract the background (original image without the person)
        # Create inverted mask (where the person is black, background is white)
        inverted_mask = Image.new('L', (width, height), 255)
        inverted_mask.paste(0, (0, 0), person_only)
        
        # Extract just the background
        background = img.copy()
        background.putalpha(inverted_mask)
        
        # Step 4: Create the text layer
        text_layer = Image.new('RGBA', (width, height), (0, 0, 0, 0))
        draw = ImageDraw.Draw(text_layer)
        
        # Prepare the text
        text = text.strip().upper()
        if not text:
            text = "SAMPLE TEXT"
            
        # Set up the font size - make it larger since we're just placing it once
        font_size = max(20, int(min(width, height) * (font_size / 100)))
        
        # Try several fonts
        font = None
        try:
            font_paths = [
                "arial.ttf", 
                "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
                "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
                "/usr/share/fonts/truetype/freefont/FreeSansBold.ttf"
            ]
            
            for font_path in font_paths:
                try:
                    font = ImageFont.truetype(font_path, font_size)
                    break
                except:
                    continue
        except:
            pass
            
        if font is None:
            font = ImageFont.load_default()
        
        # Get text dimensions
        # Try to use getbbox first for newer PIL versions
        try:
            text_width = font.getbbox(text)[2] - font.getbbox(text)[0]
            text_height = font.getbbox(text)[3] - font.getbbox(text)[1]
        except:
            # Fallback for older PIL versions
            try:
                text_width = font.getsize(text)[0]
                text_height = font.getsize(text)[1]
            except:
                # Second fallback - estimate size
                text_width = len(text) * (font_size * 0.6)
                text_height = font_size * 1.2
        
        # Position in the upper middle section
        # Divide the image into a 3x3 grid and put the text in the top-center cell
        grid_width = width // 3
        grid_height = height // 3
        
        # Center horizontally in the middle third
        x_center = width // 2 - text_width // 2
        
        # Position vertically in the top third
        y_top = grid_height // 2 - text_height // 2
        
        # Draw the text with specified opacity and color
        draw.text((x_center, y_top), text, font=font, fill=text_color_rgb + (int(text_opacity * 255),))
        
        # Step 5: Composite all layers together
        # First, composite the text on top of the background
        background_with_text = Image.alpha_composite(background, text_layer)
        
        # Then, composite the person on top
        final_image = Image.alpha_composite(background_with_text, person_mask)
        
        # Convert to RGB for display
        return np.array(final_image.convert('RGB'))
    
    except Exception as e:
        print(f"Error processing image: {e}")
        return input_image  # Return original image on error

# Create Gradio interface
with gr.Blocks(title="Text Behind Image") as demo:
    gr.Markdown("# Text Behind Image")
    gr.Markdown("Upload an image with a person and add text behind them")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Upload Image", type="numpy")
            text_input = gr.Textbox(label="Text to place behind", placeholder="Enter text here...")
            
            with gr.Row():
                text_color = gr.ColorPicker(label="Text Color", value="#000000")
                font_size = gr.Slider(label="Font Size (%)", minimum=1, maximum=30, value=10, step=1)
                text_opacity = gr.Slider(label="Text Opacity", minimum=0.1, maximum=1.0, value=0.8, step=0.1)
            
            submit_btn = gr.Button("Generate", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(label="Result", type="numpy")
    
    submit_btn.click(
        fn=text_behind_image,
        inputs=[input_image, text_input, text_color, font_size, text_opacity],
        outputs=output_image
    )
    
    gr.Markdown("## How it works")
    gr.Markdown("1. Upload an image with a person")
    gr.Markdown("2. Enter the text you want to place behind the person")
    gr.Markdown("3. Customize text color, size, and opacity")
    gr.Markdown("4. Click 'Generate' to create your image")

# Launch the app
demo.launch()