import gradio as gr
import cv2
import numpy as np
from PIL import Image
import os
import zipfile
import tempfile
import shutil

def process_image(img, strength, upscale=False):
    if upscale:
        new_size = (img.width * 2, img.height * 2)
        img = img.resize(new_size, resample=Image.LANCZOS)

    img_cv = np.array(img.convert('L'))
    img_cv = img_cv.astype(np.float32) / 255.0

    sobelx = cv2.Sobel(img_cv, cv2.CV_64F, 1, 0, ksize=3)
    sobely = cv2.Sobel(img_cv, cv2.CV_64F, 0, 1, ksize=3)


    nx = -sobelx * strength
    ny = -sobely * strength
    nz = np.ones_like(img_cv)


    length = np.sqrt(nx**2 + ny**2 + nz**2)
    nx /= length
    ny /= length
    nz /= length


    nx = (nx + 1) / 2 * 255
    ny = (ny + 1) / 2 * 255
    nz = (nz + 1) / 2 * 255

    normal_map = np.stack((nx, ny, nz), axis=-1).astype(np.uint8)
    return Image.fromarray(normal_map)

def process_single(image, strength, upscale):
    if image is None:
        return None
    return process_image(image, strength, upscale)

def process_files(file_objs, strength, upscale):
    if not file_objs:
        return None
    
    # Create temp directory and ZIP file
    temp_dir = tempfile.mkdtemp()
    zip_filename = "normal_maps.zip"
    zip_path = os.path.join(tempfile.gettempdir(), zip_filename)
    
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for file_obj in file_objs:
            file_path = file_obj.name if hasattr(file_obj, 'name') else file_obj
            file_name = os.path.basename(file_path)
            
            if file_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp')):
                try:
                    img = Image.open(file_path)
                    normal_map = process_image(img, strength, upscale)
                    
                    base_name, ext = os.path.splitext(file_name)
                    out_name = f"{base_name}_normalmap{ext}"
                        
                    out_path = os.path.join(temp_dir, out_name)
                    try:
                        normal_map.save(out_path)
                    except:
                        out_name = f"{base_name}_normalmap.png"
                        out_path = os.path.join(temp_dir, out_name)
                        normal_map.save(out_path, format="PNG")
                        
                    zipf.write(out_path, arcname=out_name)
                except Exception as e:
                    print(f"Error ({file_path}): {e}")
    
    # Cleanup temp directory
    shutil.rmtree(temp_dir)
    return zip_path

# Gradio Interface
with gr.Blocks(title="Normal Map Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🗺️ Normal Map Generator")
    gr.Markdown("Create normal maps from standard images for 3D modeling and texturing.")
    
    with gr.Tabs():
        # SINGLE IMAGE TAB
        with gr.Tab("Single Image"):
            with gr.Row():
                with gr.Column():
                    img_input = gr.Image(type="pil", label="Input Image")
                    strength_slider_1 = gr.Slider(minimum=0.1, maximum=10.0, value=2.0, step=0.1, label="Normal Map Strength")
                    upscale_check_1 = gr.Checkbox(label="Upscale 2x (High Quality Smoothing)", value=False)
                    btn_single = gr.Button("Generate", variant="primary")
                with gr.Column():
                    img_output = gr.Image(type="pil", label="Normal Map Result", format="png")
            
            btn_single.click(fn=process_single, inputs=[img_input, strength_slider_1, upscale_check_1], outputs=img_output)
            
        # FOLDER/BATCH TAB
        with gr.Tab("Folder (Batch Processing)"):
            gr.Markdown("Select a folder or drag and drop multiple images. All images will be processed and provided as a ZIP file.")
            with gr.Row():
                with gr.Column():
                    # file_count="directory" allows folder upload
                    folder_input = gr.File(file_count="directory", label="Upload Folder / Files")
                    strength_slider_2 = gr.Slider(minimum=0.1, maximum=10.0, value=2.0, step=0.1, label="Normal Map Strength")
                    upscale_check_2 = gr.Checkbox(label="Upscale 2x (High Quality Smoothing)", value=False)
                    btn_folder = gr.Button("Generate Batch and Download", variant="primary")
                with gr.Column():
                    file_output = gr.File(label="Download ZIP File")
            
            btn_folder.click(fn=process_files, inputs=[folder_input, strength_slider_2, upscale_check_2], outputs=file_output)

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