import gradio as gr
from PIL import Image, ImageDraw
import utils
import io
import os
import cv2
import numpy as np
import tempfile

def highlight_watermark(image, coords=(0, 0, 100, 50), color="red", width=5):
    try:
        if isinstance(image, np.ndarray):
            image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(image)
        x, y, w, h = coords
        draw.rectangle((max(0, x - 5), max(0, y - 5), x + w + 5, y + h + 5), outline=color, width=width)
        return image
    except Exception as e:
        print(f"Error highlighting: {e}")
        return image

def create_output_file(image, original_filename):
    try:
        # Create temp directory if it doesn't exist
        temp_dir = "temp"
        os.makedirs(temp_dir, exist_ok=True)
        
        # Generate output filename
        base_name = os.path.splitext(os.path.basename(original_filename))[0]
        output_path = os.path.join(temp_dir, f"{base_name}_watermarked.png")
        
        # Save image
        if isinstance(image, np.ndarray):
            image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        image.save(output_path, format='PNG')
        
        return output_path
    except Exception as e:
        print(f"Error creating output file: {e}")
        return None

def choose_encode(inp_im, inp_mark, cho):
    try:
        if not inp_im:
            return None, "Please upload an image.", None
        if not inp_mark:
            return None, "Please enter watermark text.", None

        valid_image_types = ["png", "jpg", "jpeg", "gif", "bmp", "tiff", "webp"]
        if not any(inp_im.lower().endswith("." + ext) for ext in valid_image_types):
            return None, "Invalid image format. Use PNG or JPG.", None

        if cho == "stegan":
            watermarked_image, error_msg = utils.encode(inp_im, inp_mark)
            
            if error_msg:
                return None, error_msg, None
                
            image_pil = Image.fromarray(cv2.cvtColor(watermarked_image, cv2.COLOR_BGR2RGB))
            highlighted_image = highlight_watermark(image_pil.copy(), (0, 0, image_pil.width, image_pil.height), color="blue")
            
            output_path = create_output_file(highlighted_image, inp_im)
            if not output_path:
                return None, "Error saving watermarked image.", None
            
            return highlighted_image, "Steganography watermark added successfully.", output_path

        elif cho == "pnginfo":
            out_im_path, error_msg = utils.png_encode(inp_im, inp_mark)
            if error_msg:
                return None, error_msg, None
                
            try:
                img = Image.open(out_im_path)
                coords = utils.png_encode_coords if hasattr(utils, 'png_encode_coords') else (50, 50, 200, 50)
                highlighted_image = highlight_watermark(img.copy(), coords)
                
                output_path = create_output_file(highlighted_image, inp_im)
                if not output_path:
                    return None, "Error saving watermarked image.", None
                
                return highlighted_image, "PNG metadata watermark added successfully.", output_path
            except Exception as e:
                return None, f"Error processing PNG: {e}", None

    except Exception as e:
        return None, f"An unexpected error occurred: {e}", None

def detect_watermark(det_im):
    if not det_im:
        return "Please upload an image."

    valid_image_types = ["png", "jpg", "jpeg", "gif", "bmp", "tiff", "webp"]
    if not any(det_im.lower().endswith("." + ext) for ext in valid_image_types):
        return "Invalid image format. Use PNG or JPG."

    detected_text = utils.decode(det_im)
    if not detected_text or detected_text.startswith("Error"):
        return "No watermark detected or not encoded with this tool."

    return detected_text

# Create Gradio interface
with gr.Blocks() as app:
    with gr.Tab("Add Watermark"):
        cho = gr.Radio(choices=["stegan", "pnginfo"], value="stegan", label="Watermark Method")
        with gr.Row():
            with gr.Column():
                inp_im = gr.Image(label="Input Image", type="filepath")
                inp_mark = gr.Textbox(label="Watermark Text")
                mark_btn = gr.Button("Add Watermark")
                msg_box = gr.Textbox(label="System Message")
            with gr.Column():
                out_im = gr.Image(label="Watermarked Image")
                file_output = gr.File(label="Download Watermarked Image")
    
    with gr.Tab("Detect Watermark"):
        with gr.Row():
            with gr.Column():
                det_im = gr.Image(label="Watermarked Image", type="filepath")
                det_btn = gr.Button("Detect Watermark")
            with gr.Column():
                det_msg = gr.Textbox(label="Detected Watermark", lines=6)

    mark_btn.click(choose_encode, inputs=[inp_im, inp_mark, cho], outputs=[out_im, msg_box, file_output])
    det_btn.click(detect_watermark, inputs=[det_im], outputs=[det_msg])

if __name__ == "__main__":
    app.launch()