import gradio as gr
import cv2
import numpy as np
from PIL import Image

# Function to add watermark (logo or text) to the image
def add_watermark(image, watermark_type, text=None, logo=None, opacity=1.0):
    image_logow = image.resize((500, 300))
    image_logow = np.array(image_logow.convert('RGB'))
    
    h_image, w_image, _ = image_logow.shape

    if watermark_type == "Logo" and logo is not None:
        logo = logo.resize((int(w_image * 0.3), int(h_image * 0.1)))  # Resize logo to 30% of the image width
        logo = np.array(logo.convert('RGB'))
        h_logo, w_logo, _ = logo.shape

        # Center placement for logo
        center_y = int(h_image / 2)
        center_x = int(w_image / 2)
        top_y = center_y - int(h_logo / 2)
        left_x = center_x - int(w_logo / 2)
        bottom_y = top_y + h_logo
        right_x = left_x + w_logo

        # Get ROI and add logo with opacity
        roi = image_logow[top_y: bottom_y, left_x: right_x]
        result = cv2.addWeighted(roi, 1, logo, opacity, 0)

        # Replace the ROI on the image
        image_logow[top_y: bottom_y, left_x: right_x] = result

        return Image.fromarray(image_logow)

    elif watermark_type == "Text" and text:
        # Add text watermark at the bottom-right corner
        cv2.putText(image_logow, text=text, org=(w_image - 100, h_image - 10), fontFace=cv2.FONT_HERSHEY_COMPLEX, 
                    fontScale=0.6, color=(0, 0, 255), thickness=2, lineType=cv2.LINE_AA)

        return Image.fromarray(image_logow)

    return image  # Return original image if no watermark added

# Create the Gradio interface
def gradio_interface():
    with gr.Blocks() as demo:
        gr.Markdown("## Image Watermarking App")
        
        gr.Markdown("""
        ### Add Watermarks to Your Images
        This application allows you to add custom watermarks to your images. You can either add a **Text Watermark** or use a **Logo Watermark**.
        
        **Steps**:
        1. Upload your image.
        2. If using a logo watermark, upload your logo.
        3. Select the type of watermark (text or logo).
        4. Customize the watermark by adding text or adjusting logo opacity.
        5. See the output image with the watermark applied.
        """)

        # Upload original image and logo
        image_input = gr.Image(label="Upload Image", type="pil")  # Upload original image
        logo_input = gr.Image(label="Upload Logo (optional, for Logo Watermark)", type="pil")  # Upload logo (optional)

        # Select watermark type and adjust parameters
        watermark_type = gr.Radio(["Text", "Logo"], label="Watermark Type", value="Text")
        text_input = gr.Textbox(label="Watermark Text (for Text Watermark)", value="MyArt.Inc", placeholder="Enter text for text watermark")
        opacity_slider = gr.Slider(0.1, 1.0, step=0.1, label="Logo Opacity (for Logo Watermark)", value=1.0)

        # Output watermarked image
        output_image = gr.Image(label="Watermarked Image")

        # Function to update the watermark based on user input
        def update_watermark(image, watermark_type, text, logo, opacity):
            if watermark_type == "Logo" and logo is None:
                return "Please upload a logo for the logo watermark."
            return add_watermark(image, watermark_type, text, logo, opacity)

        # Link inputs to output
        inputs = [image_input, watermark_type, text_input, logo_input, opacity_slider]
        gr.Button("Apply Watermark").click(update_watermark, inputs, output_image)

    demo.launch()

# Run the interface
gradio_interface()