import gradio as gr
import numpy as np
from PIL import Image, ImageEnhance
import onnxruntime as ort
import cv2
import uuid
import base64
from io import BytesIO
from cryptography.fernet import Fernet

def decrypt(encM):
    cipher = Fernet('kPcHGChMOB4inNMH-Xb_SgSgOxPN43jXuNEej76XkJc=')
    with open(encM, "rb") as f:
        decrypted_model = cipher.decrypt(f.read())
    return decrypted_model

class NAFNetProcessor:
    def __init__(self, model_path):
        """Initialize NAFNet processor"""
        print("Loading NAFNet model...")
        self.session = ort.InferenceSession(
            decrypt(model_path),
            providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
        )
        self.input_name = self.session.get_inputs()[0].name
        print("Model loaded successfully")
    
    def deblur_image(self, pil_image):
        """Process image with Padding Helper"""
        img_np = np.array(pil_image.convert('RGB'))
        h, w, _ = img_np.shape

        # 1. Padding Helper: Calculate padding to multiple of 8
        pad_h = (8 - h % 8) % 8
        pad_w = (8 - w % 8) % 8
        img_padded = cv2.copyMakeBorder(img_np, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT)

        # 2. Preprocess
        img_input = cv2.cvtColor(img_padded, cv2.COLOR_RGB2BGR)
        img_input = img_input.astype(np.float32) / 255.0
        img_input = img_input.transpose(2, 0, 1)[np.newaxis, ...]

        # 3. Inference
        results = self.session.run(None, {self.input_name: img_input})
        output = results[0][0].transpose(1, 2, 0)
        
        # 4. Postprocess
        output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
        output = np.clip(output * 255, 0, 255).astype(np.uint8)

        # 5. Crop Helper: Remove the padding
        output = output[:h, :w, :]
        return Image.fromarray(output)

def create_html_comparison(original, deblurred):
    """HTML Slider with Magnifying Glass Zoom Effect"""
    if original is None or deblurred is None:
        return "<div style='text-align: center; padding: 20px;'>Upload and process to see comparison</div>"
    
    uid = "id_" + str(uuid.uuid4())[:8]
    
    def pil_to_base64(img):
        buffered = BytesIO()
        img.save(buffered, format="PNG")
        return base64.b64encode(buffered.getvalue()).decode()
    
    orig_b64 = pil_to_base64(original)
    deblur_b64 = pil_to_base64(deblurred)
    w, h = original.size
    
    # Scaling for UI
    max_w = 700
    scale = min(max_w / w, 1.0)
    dw, dh = int(w * scale), int(h * scale)

    html = f"""
    <div id="comp-{uid}" style="width:{dw}px; height:{dh}px; position:relative; overflow:hidden; margin:0 auto; cursor:crosshair; border:2px solid #444; border-radius:8px; user-select:none;">
        
        <img src="data:image/png;base64,{orig_b64}" style="width:100%; height:100%; object-fit:contain; pointer-events:none;">
        
        <div id="overlay-{uid}" style="position:absolute; top:0; left:0; width:100%; height:100%; clip-path:inset(0 0 0 50%); pointer-events:none;">
            <img src="data:image/png;base64,{deblur_b64}" style="width:{dw}px; height:{dh}px; object-fit:contain;">
        </div>
        
        <div id="handle-{uid}" style="position:absolute; top:0; left:50%; width:4px; height:100%; background:red; transform:translateX(-50%); z-index:20; pointer-events:none;">
            <div style="position:absolute; top:50%; left:50%; transform:translate(-50%,-50%); width:30px; height:30px; background:red; border-radius:50%; color:white; display:flex; align-items:center; justify-content:center; box-shadow:0 0 10px rgba(0,0,0,0.5);">↔</div>
        </div>

        <div id="lens-{uid}" style="position:absolute; width:150px; height:150px; border:3px solid white; border-radius:50%; pointer-events:none; display:none; z-index:30; overflow:hidden; box-shadow: 0 0 15px rgba(0,0,0,0.5); background:black;">
            <img id="zoom-orig-{uid}" src="data:image/png;base64,{orig_b64}" style="position:absolute; width:{dw*2}px; height:{dh*2}px; max-width:none;">
            <div id="zoom-overlay-{uid}" style="position:absolute; width:100%; height:100%; clip-path:inset(0 0 0 50%);">
                <img src="data:image/png;base64,{deblur_b64}" style="position:absolute; width:{dw*2}px; height:{dh*2}px; max-width:none;">
            </div>
        </div>

        <img src="x" onerror='(function(){{
            const container = document.getElementById("comp-{uid}");
            const overlay = document.getElementById("overlay-{uid}");
            const handle = document.getElementById("handle-{uid}");
            const lens = document.getElementById("lens-{uid}");
            const zOrig = document.getElementById("zoom-orig-{uid}");
            const zOver = document.getElementById("zoom-overlay-{uid}");
            
            let active = false;
            let sliderPct = 50;

            const update = (e) => {{
                const rect = container.getBoundingClientRect();
                const x = (e.pageX || (e.touches ? e.touches[0].pageX : 0)) - rect.left;
                const y = (e.pageY || (e.touches ? e.touches[0].pageY : 0)) - rect.top;
                
                if (active) {{
                    sliderPct = Math.max(0, Math.min(100, (x / rect.width) * 100));
                    handle.style.left = sliderPct + "%";
                    overlay.style.clipPath = "inset(0 0 0 " + sliderPct + "%)";
                    zOver.style.clipPath = "inset(0 0 0 " + sliderPct + "%)";
                }}

                // Zoom Logic
                lens.style.display = "block";
                lens.style.left = (x - 75) + "px";
                lens.style.top = (y - 75) + "px";
                
                // Position internal images for 2x zoom
                const zX = -(x * 2 - 75);
                const zY = -(y * 2 - 75);
                zOrig.style.left = zX + "px";
                zOrig.style.top = zY + "px";
                zOver.querySelector("img").style.left = zX + "px";
                zOver.querySelector("img").style.top = zY + "px";
            }};

            container.addEventListener("mousedown", () => active = true);
            window.addEventListener("mouseup", () => active = false);
            container.addEventListener("mousemove", update);
            container.addEventListener("mouseleave", () => lens.style.display = "none");
            
            // Touch
            container.addEventListener("touchstart", (e) => {{ active = true; update(e); }});
            window.addEventListener("touchend", () => active = false);
            container.addEventListener("touchmove", update);
        }})();' style="display:none;">
    </div>
    """
    return html

def create_side_by_side(original, deblurred):
    """Creates a clean horizontal concatenation of the two images"""
    if original is None or deblurred is None:
        return None
    
    # Ensure both are RGB
    orig = original.convert("RGB")
    deblur = deblurred.convert("RGB")
    
    # Resize deblurred to match original exactly (just in case)
    if deblur.size != orig.size:
        deblur = deblur.resize(orig.size, Image.Resampling.LANCZOS)
    
    # Create a canvas for both + 10px gap
    dst = Image.new('RGB', (orig.width + deblur.width + 10, orig.height), (255, 255, 255))
    dst.paste(orig, (0, 0))
    dst.paste(deblur, (orig.width + 10, 0))
    
    return dst
    

def apply_adjustments(image, brightness, contrast, sharpness):
    """Adjusts a PIL image dynamically"""
    if image is None: return None
    
    # Brightness
    enhancer = ImageEnhance.Brightness(image)
    image = enhancer.enhance(brightness)
    
    # Contrast
    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(contrast)
    
    # Sharpness (The 'Pop' factor)
    enhancer = ImageEnhance.Sharpness(image)
    image = enhancer.enhance(sharpness)
    
    return image    

# ========== MAIN APP ==========

def main():
    processor = NAFNetProcessor("nafnet.onnx")

    def deblur_stage(input_img):
        if input_img is None: return None, None, None, None, "Ready"
        if input_img.size[0] < 400 or input_img.size[1] < 400:
            raise gr.Error("Minimum size is 512x512!")
        # Run AI
        deblurred_raw = processor.deblur_image(input_img)
        # Generate initial views
        html_view = create_html_comparison(input_img, deblurred_raw)
        sbs_view = create_side_by_side(input_img, deblurred_raw)
        return deblurred_raw, html_view, sbs_view, deblurred_raw, "✅ AI processing complete. Use sliders to tune!"

    def update_adjustments(input_img, deblurred_raw, b, c, s):
        if deblurred_raw is None: return None, None, None
        adjusted = apply_adjustments(deblurred_raw, b, c, s)
        return create_html_comparison(input_img, adjusted), create_side_by_side(input_img, adjusted), adjusted

    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        cached_deblurred = gr.State()
        gr.Markdown("# 🎮 NAFNet Ultra Deblur + Smart Zoom")
        
        with gr.Row():
            with gr.Column(scale=1):
                input_view = gr.Image(label="Input", type="pil")
                btn = gr.Button("🚀 Step 1: Run AI Deblur", variant="primary")

                gr.Examples(
                    examples=[
                        "./examples/ii.jpg", 
                        "./examples/plates.png",
                        "./examples/big.jpg"
                    ],
                    inputs=input_view,
                    label="Click an example to test"
                )
                with gr.Group():
                    b_slider = gr.Slider(0.5, 2.0, value=1.0, label="Brightness")
                    c_slider = gr.Slider(0.5, 2.0, value=1.0, label="Contrast")
                    s_slider = gr.Slider(0.0, 3.0, value=1.0, label="Sharpness")
            
            with gr.Column(scale=2):
                html_output = gr.HTML(label="Comparison (Hover for Zoom)")
                status_box = gr.Textbox(label="Status", interactive=False)

        with gr.Row():
            sbs_output = gr.Image(label="Side-by-Side View")
            download_result = gr.Image(label="Download Final Result")

        btn.click(deblur_stage, [input_view], [cached_deblurred, html_output, sbs_output, download_result, status_box])
        
        # Connect dynamic sliders
        sliders = [input_view, cached_deblurred, b_slider, c_slider, s_slider]
        for s in [b_slider, c_slider, s_slider]:
            s.change(update_adjustments, sliders, [html_output, sbs_output, download_result])

    return demo




    # ========== GRADIO INTERFACE ==========
    with gr.Blocks() as demo:
        gr.Markdown("""
        # 🎮 NAFNet Image Deblurring
        **Upload a blurry image and compare with deblurred result**
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                # Upload
                upload = gr.Image(
                    label="📤 Upload Image",
                    type="pil",
                    height=250
                )
                
                # Process button
                process_btn = gr.Button(
                    "🚀 Process Image",
                    variant="primary",
                    size="lg"
                )
                
                # Status
                status = gr.Textbox(
                    label="Status",
                    value="Ready to process",
                    interactive=False
                )
                
                gr.Markdown("""
                ### Instructions:
                1. Upload blurry image
                2. Click **Process Image**
                3. Wait a few seconds
                4. Drag the red line to compare
                """)
            
            with gr.Column(scale=2):
                # Interactive comparison
                html_output = gr.HTML(
                    label="🔄 Interactive Comparison",
                    value="<div style='text-align: center; padding: 50px; color: #666;'>Upload an image to begin</div>"
                )
                
                # Side-by-side
                with gr.Accordion("📸 Side-by-Side View", open=True):
                    side_by_side = gr.Image(
                        label="Original vs Deblurred",
                        type="pil",
                        height=350
                    )
        
        # Event handlers
        process_btn.click(
            fn=process_image,
            inputs=[upload],
            outputs=[html_output, side_by_side, status]
        )
        
        # Clear when new image uploaded
        upload.change(
            fn=lambda: (
                "<div style='text-align: center; padding: 50px; color: #666;'>Click 'Process Image' to start</div>",
                None,
                "Image ready for processing"
            ),
            inputs=[],
            outputs=[html_output, side_by_side, status]
        )
    
    return demo

# ========== RUN THE APP ==========
# if __name__ == "__main__":
demo = main()
demo.launch(debug=True)