import os
import gradio as gr
from zeroscratches import EraseScratches
import cv2
import numpy as np
from concurrent.futures import ThreadPoolExecutor
from PIL import Image

# ✅ Custom CSS for clean UI
custom_css = """
/* Dark theme styling */
body, .gradio-container { 
    background-color: #000000 !important; 
    color: white !important; 
/* Hide fullscreen and share buttons */
button[title="Fullscreen"], button[title="Share"] {
    display: none !important;
}    
}
/* Styling for buttons */
.gr-button { 
    background: #1e90ff !important; 
    color: white !important; 
    border: none !important; 
    padding: 10px 20px !important;   
    font-size: 14px !important; 
    width: 100% !important;  
    cursor: pointer;
}
.gr-button:hover {
    background: #0056b3 !important;
}
/* Styling for output boxes */
.gr-box, .gr-input, .gr-output { 
    background-color: #1c1c1c !important;  
    color: white !important; 
    border: 1px solid #333333 !important; 
}
/* Hide Gradio footer, branding, and header folders */
footer, header, .gradio-footer, .gradio-header, 
#footer, #header, #description, 
.svelte-1i4i8j8, .svelte-1lcdy9p, .svelte-18k01vx, 
.svelte-1qtpyo7, .svelte-1xjlq5, .svelte-1g01jrn, 
.svelte-1rs4uyx, .gradio-branding, 
#share-btn, #share, #download-btn, #download, 
.gradio-flag, .gradio-screenshot, 
.gradio-tab, .gradio-tabs, 
.gradio-tabs-wrapper, 
.gradio-header__space { 
    display: none !important; 
}
/* Remove share button in output box */
.gradio-container .gr-image .share-btn, 
.gradio-container .gr-image .share-button, 
.gradio-container .gr-image .share { 
    display: none !important; 
}
/* Make images non-draggable */
img { 
    pointer-events: none !important;  
    -webkit-user-drag: none !important; 
    user-select: none !important; 
}
"""

# ✅ Optimized image processing function
def predict(img_file):
    """Apply scratch removal with multi-threading"""
    with ThreadPoolExecutor(max_workers=4) as executor:
        future = executor.submit(process_image, img_file)
        restored_img = future.result()
    return restored_img

# ✅ Faster image loading and processing
def process_image(img_file):
    """Efficient image loading and restoration"""
    
    # Load image from file path using OpenCV
    img_path = img_file.name  
    img = cv2.imread(img_path)  

    # Check if image was properly loaded
    if img is None:
        raise ValueError("Failed to load image")

    # Convert to RGB
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # Resize large images for faster processing
    max_size = (1024, 1024)
    img = cv2.resize(img, max_size, interpolation=cv2.INTER_AREA)

    # Convert to PIL format for compatibility with zeroscratches
    img_pil = Image.fromarray(img)

    # Perform scratch removal
    restorer = EraseScratches()
    restored_img = restorer.erase(img_pil)

    return restored_img

# ✅ Gradio Interface
demo = gr.Interface(
    fn=predict,
    inputs=[gr.File(label="Upload Image")],
    outputs=[gr.Image(type="pil", label="Restored Image")],
    title="Restore",
    css=custom_css
)

# ✅ Launch Gradio app
demo.queue().launch(debug=True)