from PIL import Image
import numpy as np
import gradio as gr
import os
import tempfile

def perform_ela_with_red_dots(input_image, quality=90, scale=10, threshold=50):
    """
    Perform Error Level Analysis (ELA) on an image and overlay red dots on tampered areas.
    
    Parameters:
    - input_image: PIL Image object (input image)
    - quality: JPEG quality for recompression (default: 90)
    - scale: Factor to amplify differences for visibility (default: 10)
    - threshold: Threshold for highlighting significant differences (default: 50)
    
    Returns:
    - result_image: PIL Image object with ELA red dots overlaid
    """
    try:
        # Step 1: Convert input image to RGB
        original = input_image.convert('RGB')
        original_array = np.array(original)

        # Step 2: Save the image at a slightly lower quality (recompress)
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
            temp_path = temp_file.name
            original.save(temp_path, "JPEG", quality=quality)

        # Step 3: Load the recompressed image
        recompressed = Image.open(temp_path).convert('RGB')
        recompressed_array = np.array(recompressed)

        # Step 4: Compute the absolute difference between original and recompressed
        diff = np.abs(original_array - recompressed_array)

        # Step 5: Amplify the differences for better visibility
        ela = diff * scale
        ela = np.clip(ela, 0, 255).astype(np.uint8)

        # Step 6: Create a red-dot overlay for significant differences
        red_dots = np.zeros_like(original_array)
        mask = ela > threshold  # Highlight areas above the threshold
        red_dots[mask] = [255, 0, 0]  # Red color for significant differences

        # Step 7: Blend the red dots with the original image
        alpha = 0.5  # Transparency for blending
        blended = original_array.copy()
        blended = blended * (1 - alpha) + red_dots * alpha
        blended = blended.astype(np.uint8)

        # Step 8: Convert the result back to a PIL Image
        result_image = Image.fromarray(blended)

        # Step 9: Clean up the temporary file
        if os.path.exists(temp_path):
            os.remove(temp_path)

        return result_image

    except Exception as e:
        raise Exception(f"Error performing ELA: {e}")

# Gradio interface function
def ela_interface(input_image):
    """
    Gradio interface to process an image and return the ELA result.
    
    Parameters:
    - input_image: Input image uploaded via Gradio
    
    Returns:
    - result_image: Processed image with ELA red dots
    """
    if input_image is None:
        raise gr.Error("Please upload an image to process.")
    
    # Perform ELA on the uploaded image
    result_image = perform_ela_with_red_dots(input_image, quality=90, scale=10, threshold=50)
    return result_image

# Define the Gradio interface
interface = gr.Interface(
    fn=ela_interface,  # Function to call
    inputs=gr.Image(type="pil", label="Upload Image (JPEG)"),  # Input: Image upload
    outputs=gr.Image(type="pil", label="ELA Result with Red Dots"),  # Output: Processed image
    title="Error Level Analysis (ELA) for Image Tampering Detection",
    description="Upload a JPEG image to perform Error Level Analysis (ELA). Areas with potential tampering will be highlighted with red dots.",
    allow_flagging="never"
)

# Launch the Gradio app
if __name__ == "__main__":
    interface.launch(server_name="0.0.0.0", server_port=7860)