import gradio as gr
import os
import zipfile
import tempfile
import shutil
from PIL import Image
import io
from typing import List, Tuple, Optional
import numpy as np

class ImageResizer:
    def __init__(self):
        self.supported_formats = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp')
    
    def get_background_color(self, image: Image.Image) -> tuple:
        """
        Smart background color detection from image corners.
        Samples 10x10 pixel areas from all four corners and calculates average.
        """
        width, height = image.size
        
        # Define corner regions (10x10 pixels)
        corner_size = min(10, width // 4, height // 4)
        
        corners = [
            (0, 0, corner_size, corner_size),  # Top-left
            (width - corner_size, 0, width, corner_size),  # Top-right
            (0, height - corner_size, corner_size, height),  # Bottom-left
            (width - corner_size, height - corner_size, width, height)  # Bottom-right
        ]
        
        total_r, total_g, total_b = 0, 0, 0
        total_pixels = 0
        
        for corner in corners:
            corner_region = image.crop(corner)
            # Convert to RGB if not already
            if corner_region.mode != 'RGB':
                corner_region = corner_region.convert('RGB')
            
            # Get average color of this corner
            pixels = list(corner_region.getdata())
            for r, g, b in pixels:
                total_r += r
                total_g += g
                total_b += b
                total_pixels += 1
        
        if total_pixels > 0:
            avg_r = total_r // total_pixels
            avg_g = total_g // total_pixels
            avg_b = total_b // total_pixels
            return (avg_r, avg_g, avg_b)
        else:
            return (255, 255, 255)  # Default to white
    
    def resize_image(self, image: Image.Image, width: int, height: int, maintain_aspect: bool = True, png_bg_option: str = "auto", custom_color: tuple = None, is_png: bool = False) -> Image.Image:
        """
        Resize image using smart canvas padding instead of traditional resizing.
        """
        original_width, original_height = image.size
        
        # Calculate target dimensions
        if maintain_aspect:
            # Calculate aspect ratios
            original_aspect = original_width / original_height
            target_aspect = width / height
            
            if original_aspect > target_aspect:
                # Image is wider, fit to width
                new_width = width
                new_height = int(width / original_aspect)
            else:
                # Image is taller, fit to height
                new_height = height
                new_width = int(height * original_aspect)
        else:
            new_width = width
            new_height = height
        
        # Only resize if the image is larger than target dimensions
        if new_width < original_width or new_height < original_height:
            # Use LANCZOS for high-quality downsampling
            resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
        else:
            # Keep original size if it's smaller than target
            resized_image = image.copy()
            new_width, new_height = resized_image.size
        
        # Create canvas with padding
        canvas = Image.new('RGB', (width, height), (255, 255, 255))
        
        # Handle PNG background color
        if is_png or (hasattr(image, 'format') and image.format == 'PNG'):
            if png_bg_option == "auto":
                bg_color = self.get_background_color(image)
            elif png_bg_option == "black":
                bg_color = (0, 0, 0)
            elif png_bg_option == "white":
                bg_color = (255, 255, 255)
            elif png_bg_option == "custom" and custom_color:
                bg_color = custom_color
            else:
                bg_color = (255, 255, 255)
            
            canvas = Image.new('RGB', (width, height), bg_color)
        
        # Calculate position to center the image
        x = (width - new_width) // 2
        y = (height - new_height) // 2
        
        # Paste the resized image onto the canvas
        if resized_image.mode == 'RGBA':
            canvas.paste(resized_image, (x, y), resized_image)
        else:
            canvas.paste(resized_image, (x, y))
        
        return canvas
    
    def hex_to_rgb(self, hex_color: str) -> tuple:
        """Convert hex color to RGB tuple."""
        hex_color = hex_color.lstrip('#')
        return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))

def process_single_image(image, width, height, maintain_aspect, png_bg_option, custom_color_hex):
    """Process a single image with Gradio interface."""
    if image is None:
        return None, "Please upload an image first."
    
    try:
        # Convert Gradio image to PIL Image
        if isinstance(image, np.ndarray):
            pil_image = Image.fromarray(image)
        else:
            pil_image = image
        
        # Validate dimensions
        if width <= 0 or height <= 0:
            return None, "Width and height must be positive numbers."
        
        # Initialize resizer
        resizer = ImageResizer()
        
        # Handle custom color
        custom_color = None
        if png_bg_option == "custom":
            try:
                custom_color = resizer.hex_to_rgb(custom_color_hex)
            except:
                custom_color = (255, 255, 255)  # Default to white if invalid
        
        # Check if image is PNG
        is_png = hasattr(pil_image, 'format') and pil_image.format == 'PNG'
        
        # Resize image
        resized_image = resizer.resize_image(
            pil_image, width, height, maintain_aspect, 
            png_bg_option, custom_color, is_png
        )
        
        return resized_image, f"✅ Image resized successfully to {width}x{height}!"
        
    except Exception as e:
        return None, f"❌ Error processing image: {str(e)}"

def process_folder_images(files, width, height, maintain_aspect, png_bg_option, custom_color_hex):
    """Process multiple images from folder upload."""
    if not files:
        return [], "Please upload image files first."
    
    try:
        # Validate dimensions
        if width <= 0 or height <= 0:
            return [], "Width and height must be positive numbers."
        
        # Initialize resizer
        resizer = ImageResizer()
        
        # Handle custom color
        custom_color = None
        if png_bg_option == "custom":
            try:
                custom_color = resizer.hex_to_rgb(custom_color_hex)
            except:
                custom_color = (255, 255, 255)  # Default to white if invalid
        
        processed_images = []
        processed_count = 0
        
        for file in files:
            try:
                # Open image
                image = Image.open(file.name)
                
                # Check if image is PNG
                is_png = file.name.lower().endswith('.png')
                
                # Resize image
                resized_image = resizer.resize_image(
                    image, width, height, maintain_aspect, 
                    png_bg_option, custom_color, is_png
                )
                
                processed_images.append(resized_image)
                processed_count += 1
                
            except Exception as e:
                print(f"Error processing {file.name}: {str(e)}")
                continue
        
        status_msg = f"✅ Successfully processed {processed_count} out of {len(files)} images to {width}x{height}!"
        return processed_images, status_msg
        
    except Exception as e:
        return [], f"❌ Error processing folder: {str(e)}"

def process_zip_file(zip_file, width, height, maintain_aspect, png_bg_option, custom_color_hex):
    """Process images from uploaded ZIP file."""
    if zip_file is None:
        return [], "Please upload a ZIP file first."
    
    try:
        # Validate dimensions
        if width <= 0 or height <= 0:
            return [], "Width and height must be positive numbers."
        
        # Initialize resizer
        resizer = ImageResizer()
        
        # Handle custom color
        custom_color = None
        if png_bg_option == "custom":
            try:
                custom_color = resizer.hex_to_rgb(custom_color_hex)
            except:
                custom_color = (255, 255, 255)  # Default to white if invalid
        
        processed_images = []
        processed_count = 0
        
        # Create temporary directory
        with tempfile.TemporaryDirectory() as temp_dir:
            # Extract ZIP file
            with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
                zip_ref.extractall(temp_dir)
            
            # Process all images in the extracted folder
            for root, dirs, files in os.walk(temp_dir):
                for file in files:
                    if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp')):
                        try:
                            file_path = os.path.join(root, file)
                            image = Image.open(file_path)
                            
                            # Check if image is PNG
                            is_png = file.lower().endswith('.png')
                            
                            # Resize image
                            resized_image = resizer.resize_image(
                                image, width, height, maintain_aspect, 
                                png_bg_option, custom_color, is_png
                            )
                            
                            processed_images.append(resized_image)
                            processed_count += 1
                            
                        except Exception as e:
                            print(f"Error processing {file}: {str(e)}")
                            continue
        
        status_msg = f"✅ Successfully processed {processed_count} images from ZIP file to {width}x{height}!"
        return processed_images, status_msg
        
    except Exception as e:
        return [], f"❌ Error processing ZIP file: {str(e)}"

# Create Gradio interface
def create_gradio_app():
    with gr.Blocks(title="🖼️ Image Resizer Pro - Gradio Edition", theme=gr.themes.Soft()) as app:
        gr.Markdown("# 🖼️ Image Resizer Pro - Gradio Edition")
        gr.Markdown("**Smart Canvas Padding & PNG Background Selector** - Resize images with intelligent padding instead of stretching")
        
        with gr.Tabs():
            # Single Image Tab
            with gr.TabItem("📷 Single Image"):
                with gr.Row():
                    with gr.Column(scale=1):
                        single_image_input = gr.Image(type="pil", label="Upload Image")
                        
                        with gr.Row():
                            single_width = gr.Number(value=800, label="Width", minimum=1, maximum=10000)
                            single_height = gr.Number(value=600, label="Height", minimum=1, maximum=10000)
                        
                        single_maintain_aspect = gr.Checkbox(value=True, label="Maintain aspect ratio")
                        
                        with gr.Group():
                            gr.Markdown("**PNG Background Color**")
                            single_png_bg = gr.Radio(
                                choices=["auto", "black", "white", "custom"],
                                value="auto",
                                label="Background Option",
                                info="Choose background color for PNG images"
                            )
                            single_custom_color = gr.Textbox(
                                value="#FFFFFF",
                                label="Custom Color (Hex)",
                                placeholder="#FFFFFF",
                                visible=False
                            )
                        
                        single_process_btn = gr.Button("🔄 Resize Image", variant="primary")
                    
                    with gr.Column(scale=1):
                        single_output = gr.Image(label="Resized Image")
                        single_status = gr.Textbox(label="Status", interactive=False)
                
                # Show/hide custom color input
                single_png_bg.change(
                    lambda x: gr.update(visible=(x == "custom")),
                    inputs=[single_png_bg],
                    outputs=[single_custom_color]
                )
                
                # Process single image
                single_process_btn.click(
                    process_single_image,
                    inputs=[single_image_input, single_width, single_height, single_maintain_aspect, single_png_bg, single_custom_color],
                    outputs=[single_output, single_status]
                )
            
            # Folder Processing Tab
            with gr.TabItem("📁 Folder Processing"):
                with gr.Row():
                    with gr.Column(scale=1):
                        folder_files_input = gr.File(
                            file_count="multiple",
                            file_types=["image"],
                            label="Upload Multiple Images"
                        )
                        
                        with gr.Row():
                            folder_width = gr.Number(value=800, label="Width", minimum=1, maximum=10000)
                            folder_height = gr.Number(value=600, label="Height", minimum=1, maximum=10000)
                        
                        folder_maintain_aspect = gr.Checkbox(value=True, label="Maintain aspect ratio")
                        
                        with gr.Group():
                            gr.Markdown("**PNG Background Color**")
                            folder_png_bg = gr.Radio(
                                choices=["auto", "black", "white", "custom"],
                                value="auto",
                                label="Background Option",
                                info="Choose background color for PNG images"
                            )
                            folder_custom_color = gr.Textbox(
                                value="#FFFFFF",
                                label="Custom Color (Hex)",
                                placeholder="#FFFFFF",
                                visible=False
                            )
                        
                        folder_process_btn = gr.Button("🔄 Process Folder", variant="primary")
                    
                    with gr.Column(scale=1):
                        folder_output = gr.Gallery(label="Processed Images", columns=3, rows=2)
                        folder_status = gr.Textbox(label="Status", interactive=False)
                
                # Show/hide custom color input
                folder_png_bg.change(
                    lambda x: gr.update(visible=(x == "custom")),
                    inputs=[folder_png_bg],
                    outputs=[folder_custom_color]
                )
                
                # Process folder
                folder_process_btn.click(
                    process_folder_images,
                    inputs=[folder_files_input, folder_width, folder_height, folder_maintain_aspect, folder_png_bg, folder_custom_color],
                    outputs=[folder_output, folder_status]
                )
            
            # ZIP Processing Tab
            with gr.TabItem("📦 ZIP Processing"):
                with gr.Row():
                    with gr.Column(scale=1):
                        zip_file_input = gr.File(
                            file_types=[".zip"],
                            label="Upload ZIP File"
                        )
                        
                        with gr.Row():
                            zip_width = gr.Number(value=800, label="Width", minimum=1, maximum=10000)
                            zip_height = gr.Number(value=600, label="Height", minimum=1, maximum=10000)
                        
                        zip_maintain_aspect = gr.Checkbox(value=True, label="Maintain aspect ratio")
                        
                        with gr.Group():
                            gr.Markdown("**PNG Background Color**")
                            zip_png_bg = gr.Radio(
                                choices=["auto", "black", "white", "custom"],
                                value="auto",
                                label="Background Option",
                                info="Choose background color for PNG images"
                            )
                            zip_custom_color = gr.Textbox(
                                value="#FFFFFF",
                                label="Custom Color (Hex)",
                                placeholder="#FFFFFF",
                                visible=False
                            )
                        
                        zip_process_btn = gr.Button("🔄 Process ZIP", variant="primary")
                    
                    with gr.Column(scale=1):
                        zip_output = gr.Gallery(label="Processed Images", columns=3, rows=2)
                        zip_status = gr.Textbox(label="Status", interactive=False)
                
                # Show/hide custom color input
                zip_png_bg.change(
                    lambda x: gr.update(visible=(x == "custom")),
                    inputs=[zip_png_bg],
                    outputs=[zip_custom_color]
                )
                
                # Process ZIP
                zip_process_btn.click(
                    process_zip_file,
                    inputs=[zip_file_input, zip_width, zip_height, zip_maintain_aspect, zip_png_bg, zip_custom_color],
                    outputs=[zip_output, zip_status]
                )
        
        # Footer
        gr.Markdown("---")
        gr.Markdown("**🖼️ Image Resizer Pro v3.2** - Smart Canvas Padding & PNG Background Selector Edition")
        gr.Markdown("Features: Smart background detection, Canvas padding, Multi-format support (JPG, PNG, BMP, TIFF, WEBP)")
    
    return app

if __name__ == "__main__":
    app = create_gradio_app()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True
    )