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 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)
    
    # Using min-width for better responsive behavior
    # sizes_str = "(min-width: 1000px) 45vw, (min-width: 780px) 696px, (min-width: 580px) 480px, calc(100vw - 60px)"
    
    # 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 get_path_components(filepath):
    """Split filepath into directory and filename"""
    dirpath = os.path.dirname(filepath)
    filename = os.path.basename(filepath)
    base_filename = os.path.splitext(filename)[0]
    return dirpath, base_filename

def process_image(input_img, input_path):
    """Main processing function"""
    if input_img 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:
        # Get path components
        if input_path and input_path.strip():
            dirpath, base_filename = get_path_components(input_path.strip())
        else:
            # If no path provided, use default
            dirpath, base_filename = "", "image"
            
        # Convert to PIL Image if needed
        if isinstance(input_img, np.ndarray):
            img = Image.fromarray(input_img)
        else:
            img = input_img
            
        # Convert to RGB if necessary
        if img.mode != 'RGB':
            img = img.convert('RGB')
        
        # Process each size
        for width, height in SIZES:
            # For zip file, only use filename without path
            output_path = os.path.join(temp_dir, f"{base_filename}-{width}x{height}.jpg")
            resized = resize_image(img, (width, height))
            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 full path
        full_path = dirpath + ("/" if dirpath else "") + base_filename
        html_snippet = generate_html_snippet(full_path)
        
        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
    
    Optional: Specify the full path where images will be stored (e.g., 'assets/images/about/mock-up.png')
    """)
    
    with gr.Row():
        with gr.Column():
            with gr.Row():
                input_image = gr.Image(
                    label="Upload Original Image",
                    type="pil",
                    show_label=True
                )
            with gr.Row():
                input_path = gr.Textbox(
                    label="Image Path (optional)",
                    placeholder="e.g., assets/images/about/mock-up.png",
                    value=""
                )
            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")
    
    def handle_upload(img):
        """Handle image upload to get filename"""
        if isinstance(img, dict) and 'name' in img:
            return img['name']
        return ""
    
    # Update filename when image is uploaded
    input_image.upload(
        fn=handle_upload,
        inputs=[input_image],
        outputs=[input_path]
    )
    
    # Process button click
    process_btn.click(
        fn=process_image,
        inputs=[input_image, input_path],
        outputs=[output_zip, output_html, output_message]
    )

# Launch the app with share=True for public access
app.launch(share=True)