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

# Define target sizes for different viewports
SIZES = [
    (480, 423),
    (696, 613),
    (960, 846),
    (1095, 965),
    (1280, 1128),
    (1440, 1269),
    (1670, 1472),
    (1870, 1648),
    (2048, 1805)
]

def resize_image(img, target_size):
    """Resize image maintaining aspect ratio"""
    img_copy = img.copy()
    return img_copy.resize(target_size, PIL.Image.Resampling.LANCZOS)

def get_base_filename(original_path):
    """Extract original filename without extension"""
    base_name = os.path.basename(original_path)
    return os.path.splitext(base_name)[0]

def generate_html_snippet(image_path_prefix):
    """Generate HTML code snippet with srcset"""
    srcset_items = [f"{image_path_prefix}-{w}x{h}.jpg {w}w" for w, h in SIZES]
    srcset_str = ",\n        ".join(srcset_items)
    
    sizes_items = [
        f"(max-width: {w}px) {w}px" for w, _ in SIZES[:-1]
    ] + [f"{SIZES[-1][0]}px"]
    sizes_str = ",\n        ".join(sizes_items)
    
    # Use the middle size as default src
    default_size = SIZES[4]  # 1280x1128
    
    html = f'''<!-- Responsive Image -->
<img 
    src="{image_path_prefix}-{default_size[0]}x{default_size[1]}.jpg"
    srcset="
        {srcset_str}
    "
    sizes="
        {sizes_str}
    "
    alt="Your image description"
    style="max-width: 100%; height: auto;"
>'''
    
    return html

def process_image(input_image, original_filename):
    """Main processing function"""
    if input_image is None:
        return None, None, "Please upload an image"
    
    # Create temporary directory for processed images
    temp_dir = tempfile.mkdtemp()
    zip_path = os.path.join(temp_dir, "responsive_images.zip")
    processed_paths = []  # Initialize the list outside try block
    
    try:
        # Convert numpy array to PIL Image
        img = Image.fromarray(input_image)
        
        # Convert to RGB if necessary
        if img.mode != 'RGB':
            img = img.convert('RGB')
        
        # Get base filename without extension
        base_filename = os.path.splitext(original_filename)[0]
        
        # Process each size
        for width, height in SIZES:
            resized = resize_image(img, (width, height))
            output_path = os.path.join(temp_dir, f"{base_filename}-{width}x{height}.jpg")
            resized.save(output_path, "JPEG", quality=90)
            processed_paths.append(output_path)
        
        # Create zip file
        with zipfile.ZipFile(zip_path, 'w') as zf:
            for path in processed_paths:
                zf.write(path, os.path.basename(path))
        
        # Generate HTML snippet using original filename
        html_snippet = generate_html_snippet(f"images/{base_filename}")
        
        return zip_path, html_snippet, "Processing completed successfully!"
        
    except Exception as e:
        if os.path.exists(zip_path):
            try:
                os.remove(zip_path)
            except:
                pass
        return None, None, f"Error processing image: {str(e)}"
    
    finally:
        # Clean up temporary files except zip
        for path in processed_paths:
            try:
                os.remove(path)
            except:
                pass
        # Try to remove the temp directory if it's empty
        try:
            os.rmdir(temp_dir)
        except:
            pass

# Create Gradio interface
with gr.Blocks(title="Responsive Image Generator") as app:
    gr.Markdown("""
    # Responsive Image Generator
    
    Upload an image to generate optimized versions for different viewport sizes. 
    You'll receive:
    1. A ZIP file containing all sized versions
    2. HTML code snippet with proper srcset attributes
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="Upload Original Image",
                type="numpy"  # Using numpy type for better compatibility
            )
            process_btn = gr.Button("Process Image")
        
        with gr.Column():
            output_zip = gr.File(label="Download Processed Images")
            output_html = gr.Code(
                label="HTML Code Snippet", 
                language="html"
            )
            output_message = gr.Textbox(label="Status")
    
    process_btn.click(
        fn=process_image,
        inputs=[input_image, input_image],  # Pass the image twice to get both data and filename
        outputs=[output_zip, output_html, output_message]
    )

app.launch()