import cv2
import numpy as np
import os
import zipfile
import uuid
import gradio as gr
import uuid


def remove_watermark_area(original_image, text_mask_path):
    # Ensure the mask is binary
    text_mask = cv2.imread(text_mask_path, cv2.IMREAD_GRAYSCALE)
    _, binary_mask = cv2.threshold(text_mask, 1, 255, cv2.THRESH_BINARY)

    # Resize the mask to match the size of the original image area
    mask_resized = cv2.resize(binary_mask, (original_image.shape[1], original_image.shape[0]))

    # Expand the mask to cover more area if needed
    kernel = np.ones((5, 5), np.uint8)
    expanded_mask = cv2.dilate(mask_resized, kernel, iterations=1)

    # Inpainting using the mask
    inpainted_image = cv2.inpaint(original_image, expanded_mask, inpaintRadius=5, flags=cv2.INPAINT_TELEA)

    # Optionally apply post-processing to improve results
    cleaned_image = cv2.GaussianBlur(inpainted_image, (3, 3), 0)

    return cleaned_image
from PIL import Image

def remove_watermark(image_path,saved_path):
    if isinstance(image_path, str) and os.path.isfile(image_path):
        # Load the image using OpenCV
        image = cv2.imread(image_path) 
    elif isinstance(image_path, np.ndarray):
        # Directly use OpenCV image (NumPy array)
        image = image_path
        if len(image_path.shape) == 3 and image_path.shape[2] == 3:
            # Assuming it's in RGB format; convert to BGR
            image = cv2.cvtColor(image_path, cv2.COLOR_RGB2BGR)
        else:
            # Otherwise, assume it's already in BGR format
            image = image_path
    else:
        raise TypeError("Invalid image_path format")
    print(type(image))
    cv2.imwrite("test.jpg",image)
    image=cv2.resize(image,(1280,1280))
    # Define the area of the watermark (adjust this based on the watermark size)
    height, width, _ = image.shape
    watermark_width = 185  # Adjust based on your watermark size
    watermark_height = 185  # Adjust based on your watermark size
    x_start = 50
    y_start = height - watermark_height+17
    x_end = watermark_width-17
    y_end = height-50

    # Extract the watermark area
    watermark_area = image[y_start:y_end, x_start:x_end]
    # cv2.imwrite('watermark_area.jpg', watermark_area)

    # Create the mask for the watermark area
    text_mask_path = 'watermark_mask.png'
    cleaned_image = remove_watermark_area(watermark_area, text_mask_path)
    # cv2.imwrite('cleaned_watermark.jpg', cleaned_image)
    # Paste back the cleaned watermark on the original image
    image[y_start:y_end, x_start:x_end] = cleaned_image
    cv2.imwrite(saved_path, image)
    return image

def make_zip(image_list):
    zip_path = f"./temp/{uuid.uuid4().hex[:6]}.zip"
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for image in image_list:
            zipf.write(image, os.path.basename(image))
    return zip_path

def random_image_name():
    """Generate a random image name."""
    return str(uuid.uuid4())[:8]

    
def process_file(pil_image):
    saved_path = f"./temp/{random_image_name()}.jpg"
    remove_watermark(pil_image, saved_path)
    return saved_path, saved_path


def process_files(image_files):
    image_list = []
    if len(image_files) == 1:
        # saved_path = os.path.basename(image_files[0])
        # saved_path = f"./temp/{saved_path}"
        saved_path = f"./temp/{random_image_name()}.jpg"
        remove_watermark(image_files[0], saved_path)
        return saved_path, saved_path
    else:
        for image_path in image_files:
            # saved_path = os.path.basename(image_path)
            # saved_path = f"./temp/{saved_path}"
            saved_path = f"./temp/{random_image_name()}.jpg"
            remove_watermark(image_path, saved_path)
            image_list.append(saved_path)
        zip_path = make_zip(image_list)
        return zip_path,None


if not os.path.exists("./temp"):
    os.mkdir("./temp")


meta_examples = ["./images/1.jpg", "./images/2.jpg", "./images/3.jpg", "./images/4.jpg", "./images/5.jpg", "./images/6.jpg"]

gradio_input=[gr.Image(label='Upload an Image')]
gradio_Output=[gr.File(label='Download Image'),gr.Image(label='Display Image')]
gradio_interface = gr.Interface(fn=process_file, inputs=gradio_input,outputs=gradio_Output ,
                              title="Meta Watermark Remover",
                              examples=meta_examples)
# gradio_interface.launch(debug=True)



gradio_multiple_images = gr.Interface(
    process_files,
    [gr.File(type='filepath', file_count='multiple',label='Upload Images')],
    [gr.File(label='Download File'),gr.Image(label='Display Image')],
    cache_examples=True
)

demo = gr.TabbedInterface([gradio_interface, gradio_multiple_images], ["Meta Watermark Remover","Bluk Meta Watermark Remover"],title="Meta Watermark Remover")
demo.launch(debug=True)