import gradio as gr
from PIL import Image
import io
import time
import tempfile
from pathlib import Path
import zipfile

def compress_image(file, quality=75, format_type="webp"):
    """Simple image compression function"""
    if file is None:
        return None, "Please upload an image"
    
    try:
        # Open image
        img = Image.open(file.name)
        original_size = Path(file.name).stat().st_size / 1024
        
        # Compress to selected format
        output = io.BytesIO()
        
        if format_type == "webp":
            img.save(output, format='WEBP', quality=quality, method=4)
            extension = ".webp"
        elif format_type == "jpeg":
            if img.mode in ('RGBA', 'LA', 'P'):
                # Convert to RGB for JPEG
                rgb_img = Image.new('RGB', img.size, (255, 255, 255))
                rgb_img.paste(img, mask=img.split()[-1] if img.mode == 'RGBA' else None)
                img = rgb_img
            img.save(output, format='JPEG', quality=quality, optimize=True)
            extension = ".jpg"
        else:  # png
            img.save(output, format='PNG', optimize=True)
            extension = ".png"
        
        output.seek(0)
        compressed_size = len(output.getvalue()) / 1024
        saved_percent = (1 - compressed_size / original_size) * 100
        
        # Save to temp file for download
        temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=extension)
        temp_path.write(output.getvalue())
        temp_path.close()
        
        info = f"""✅ Success!
        
Original: {original_size:.1f} KB
Compressed: {compressed_size:.1f} KB
Saved: {saved_percent:.1f}%
Format: {format_type.upper()}
Quality: {quality}"""
        
        return temp_path.name, info
        
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

def compress_batch(files, quality=75, format_type="webp"):
    """Batch compression"""
    if not files:
        return None, "Please upload files"
    
    compressed_paths = []
    total_saved = 0
    
    with tempfile.TemporaryDirectory() as temp_dir:
        for file in files:
            try:
                img = Image.open(file.name)
                original_size = Path(file.name).stat().st_size / 1024
                
                output = io.BytesIO()
                if format_type == "webp":
                    img.save(output, format='WEBP', quality=quality)
                    ext = ".webp"
                elif format_type == "jpeg":
                    if img.mode in ('RGBA', 'LA', 'P'):
                        rgb_img = Image.new('RGB', img.size, (255, 255, 255))
                        rgb_img.paste(img, mask=img.split()[-1] if img.mode == 'RGBA' else None)
                        img = rgb_img
                    img.save(output, format='JPEG', quality=quality)
                    ext = ".jpg"
                else:
                    img.save(output, format='PNG', optimize=True)
                    ext = ".png"
                
                output.seek(0)
                compressed_size = len(output.getvalue()) / 1024
                saved = (1 - compressed_size / original_size) * 100
                total_saved += saved
                
                # Save file
                temp_file = Path(temp_dir) / f"compressed_{Path(file.name).stem}{ext}"
                with open(temp_file, 'wb') as f:
                    f.write(output.getvalue())
                compressed_paths.append(temp_file)
                
            except Exception as e:
                print(f"Error: {e}")
        
        if not compressed_paths:
            return None, "No files could be compressed"
        
        # Create ZIP
        zip_path = Path(temp_dir) / "compressed_images.zip"
        with zipfile.ZipFile(zip_path, 'w') as zipf:
            for path in compressed_paths:
                zipf.write(path, path.name)
        
        # Copy to persistent location
        final_zip = Path("/tmp") / f"batch_{int(time.time())}.zip"
        import shutil
        shutil.copy(zip_path, final_zip)
        
        avg_saved = total_saved / len(compressed_paths)
        info = f"✅ Compressed {len(compressed_paths)} files\nAverage saving: {avg_saved:.1f}%"
        
        return str(final_zip), info

# Create interface
with gr.Blocks(title="Image Compressor", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🚀 Image Compressor
    
    **Simple, fast, and private - all processing happens in your browser!**
    
    ### Features:
    - Compress JPEG, PNG, WebP
    - Batch processing with ZIP download
    - Adjustable quality
    - No uploads to server
    """)
    
    with gr.Tabs():
        with gr.TabItem("Single Image"):
            with gr.Row():
                with gr.Column():
                    input_file = gr.File(label="Upload Image", file_types=["image"])
                    format_choice = gr.Radio(
                        choices=["webp", "jpeg", "png"],
                        value="webp",
                        label="Output Format"
                    )
                    quality_slider = gr.Slider(
                        minimum=30, 
                        maximum=100, 
                        value=75, 
                        label="Quality (higher = better quality, larger file)"
                    )
                    compress_btn = gr.Button("Compress", variant="primary")
                
                with gr.Column():
                    output_file = gr.File(label="Download Compressed Image")
                    info_text = gr.Textbox(label="Results", lines=8)
        
        with gr.TabItem("Batch Processing"):
            batch_files = gr.File(
                label="Upload Multiple Images", 
                file_types=["image"], 
                file_count="multiple"
            )
            batch_format = gr.Radio(
                choices=["webp", "jpeg", "png"],
                value="webp",
                label="Output Format"
            )
            batch_quality = gr.Slider(
                minimum=30, 
                maximum=100, 
                value=75, 
                label="Quality"
            )
            batch_btn = gr.Button("Compress All", variant="primary")
            batch_output = gr.File(label="Download ZIP")
            batch_info = gr.Textbox(label="Results", lines=5)
    
    # Connect functions
    compress_btn.click(
        compress_image,
        inputs=[input_file, quality_slider, format_choice],
        outputs=[output_file, info_text]
    )
    
    batch_btn.click(
        compress_batch,
        inputs=[batch_files, batch_quality, batch_format],
        outputs=[batch_output, batch_info]
    )

# Launch
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)