#!/usr/bin/env python3
"""
png2gif.py – Turn a directory of .png frames into an animated .gif
Usage:
    python png2gif.py --src ./frames --out animation.gif --fps 12
"""
import argparse
from pathlib import Path
from PIL import Image
import re
import gradio as gr
import os
import tempfile
import zipfile
import shutil
import datetime

def natural_key(s):
    """Sort 'frame_2.png' before 'frame_10.png'."""
    return [int(t) if t.isdigit() else t.lower() for t in re.split(r'(\d+)', s.name)]

def gather_frames(src_dir):
    pngs = sorted(Path(src_dir).glob("*.png"), key=natural_key)
    if not pngs:
        return []
    return pngs

def build_gif(pngs, out_file, fps, loop):
    ms_per_frame = int(1000 / fps)
    frames = [Image.open(png).convert("RGBA") for png in pngs]

    # Warn if frame sizes differ
    w, h = frames[0].size
    if any(im.size != (w, h) for im in frames):
        print("⚠️  Frame dimensions differ; resizing to first frame’s size")
        frames = [im.resize((w, h), Image.LANCZOS) for im in frames]

    frames[0].save(
        out_file,
        save_all=True,
        append_images=frames[1:],
        duration=ms_per_frame,
        loop=loop,
        disposal=2,      # clears frame before drawing next (good for transparency)
        optimize=True,   # basic optimization
    )
    print(f"✅ GIF saved to {out_file}")

def process_batch_directories(parent_dir, fps, loop_count):
    parent_path = Path(parent_dir)
    if not parent_path.exists():
        print(f"❌ Parent directory not found: {parent_dir}")
        return
    
    if not parent_path.is_dir():
        print(f"❌ Path is not a directory: {parent_dir}")
        return
    
    # Find all subdirectories
    subdirs = [d for d in parent_path.iterdir() if d.is_dir()]
    
    if not subdirs:
        print(f"❌ No subdirectories found in {parent_dir}")
        return
    
    print(f"🔍 Found {len(subdirs)} subdirectories to process")
    
    successful = 0
    failed = 0
    
    for subdir in subdirs:
        print(f"\n📁 Processing: {subdir.name}")
        
        # Gather PNG frames from this subdirectory
        pngs = gather_frames(subdir)
        
        if not pngs:
            print(f"⚠️  No PNG files found in {subdir.name}, skipping")
            continue
        
        # Create GIF filename based on subdirectory name
        gif_name = f"{subdir.name}.gif"
        output_path = subdir / gif_name
        
        try:
            build_gif(pngs, output_path, fps, loop_count)
            print(f"✅ Created {gif_name} with {len(pngs)} frames")
            successful += 1
        except Exception as e:
            print(f"❌ Failed to create GIF for {subdir.name}: {e}")
            failed += 1
    
    # Summary
    print(f"\n📊 Batch processing complete:")
    print(f"   ✅ Successful: {successful}")
    print(f"   ❌ Failed: {failed}")
    print(f"   📁 Total processed: {successful + failed}")

def process_uploaded_files(files, fps, loop_count):
    if not files:
        return None, None, "Please upload PNG files"
    
    with tempfile.TemporaryDirectory() as temp_dir:
        png_files = []
        
        for file in files:
            if not file.name.lower().endswith('.png'):
                continue
            
            # Gradio files have a .name attribute that contains the file path
            source_path = Path(file.name)
            dest_path = Path(temp_dir) / source_path.name
            
            # Copy the file from the temporary location Gradio created
            shutil.copy2(source_path, dest_path)
            png_files.append(dest_path)
        
        if not png_files:
            return None, None, "No valid PNG files found"
        
        png_files = sorted(png_files, key=natural_key)
        
        temp_output = Path(temp_dir) / "animation.gif"
        try:
            build_gif(png_files, temp_output, fps, loop_count)
            
            # Create a persistent temporary file for Gradio to access
            persistent_output = tempfile.NamedTemporaryFile(suffix=".gif", delete=False)
            persistent_output.close()
            
            # Copy the generated GIF to the persistent location
            shutil.copy2(temp_output, persistent_output.name)
            
            # Optional: build filename for display only
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            filename_for_display = f"animation_{timestamp}.gif"
            
            return persistent_output.name, persistent_output.name, f"✅ GIF created successfully with {len(png_files)} frames\n📁 Filename: {filename_for_display}"
        except Exception as e:
            return None, None, f"❌ Error creating GIF: {str(e)}"

def process_batch_zip(zip_file, fps, loop_count):
    if not zip_file:
        return "No ZIP file uploaded", [], None
    
    with tempfile.TemporaryDirectory() as temp_dir:
        try:
            # Extract ZIP file
            zip_path = Path(zip_file.name)
            extract_dir = Path(temp_dir) / "extracted"
            
            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                zip_ref.extractall(extract_dir)
            
            # Find subdirectories containing PNG files
            subdirs = []
            for item in extract_dir.rglob("*"):
                if item.is_dir():
                    pngs = list(item.glob("*.png"))
                    if pngs:
                        subdirs.append(item)
            
            if not subdirs:
                return "❌ No subdirectories with PNG files found in ZIP", [], None
            
            progress_msg = f"🔍 Found {len(subdirs)} directories to process"
            generated_gifs = []
            results_dir = Path(temp_dir) / "results"
            results_dir.mkdir()
            
            successful = 0
            failed = 0
            
            for i, subdir in enumerate(subdirs, 1):
                progress_msg += f"\n📁 Processing {i}/{len(subdirs)}: {subdir.name}"
                
                # Gather PNG frames
                pngs = gather_frames(subdir)
                if not pngs:
                    progress_msg += f"\n⚠️  No PNG files found in {subdir.name}, skipping"
                    continue
                
                # Create GIF
                gif_name = f"{subdir.name}.gif"
                gif_path = results_dir / gif_name
                
                try:
                    build_gif(pngs, gif_path, fps, loop_count)
                    
                    # Create persistent copy for gallery display
                    persistent_gif = tempfile.NamedTemporaryFile(suffix=".gif", delete=False)
                    persistent_gif.close()
                    shutil.copy2(gif_path, persistent_gif.name)
                    generated_gifs.append(persistent_gif.name)
                    
                    progress_msg += f"\n✅ Created {gif_name} with {len(pngs)} frames"
                    successful += 1
                except Exception as e:
                    progress_msg += f"\n❌ Failed to create GIF for {subdir.name}: {e}"
                    failed += 1
            
            # Create results ZIP
            if generated_gifs:
                results_zip = tempfile.NamedTemporaryFile(suffix=".zip", delete=False)
                results_zip.close()
                
                with zipfile.ZipFile(results_zip.name, 'w') as zip_out:
                    for gif_path in results_dir.glob("*.gif"):
                        zip_out.write(gif_path, gif_path.name)
                
                progress_msg += f"\n\n📊 Batch processing complete:"
                progress_msg += f"\n   ✅ Successful: {successful}"
                progress_msg += f"\n   ❌ Failed: {failed}"
                progress_msg += f"\n   📦 Results ZIP created with {len(generated_gifs)} GIFs"
                
                return progress_msg, generated_gifs, results_zip.name
            else:
                return "❌ No GIFs were successfully created", [], None
                
        except Exception as e:
            return f"❌ Error processing ZIP file: {str(e)}", [], None

def toggle_mode(mode):
    if mode == "Single Sequence":
        return (
            gr.update(visible=True),   # single_group
            gr.update(visible=False),  # batch_group  
            gr.update(visible=True),   # single_output
            gr.update(visible=False)   # batch_output
        )
    else:
        return (
            gr.update(visible=False),  # single_group
            gr.update(visible=True),   # batch_group
            gr.update(visible=False),  # single_output  
            gr.update(visible=True)    # batch_output
        )

def process_conversion(mode, files, zip_file, fps, loop_count):
    if mode == "Single Sequence":
        if files:
            preview, download, status = process_uploaded_files(files, fps, loop_count)
            return preview, download, status, "", [], None
        else:
            return None, None, "Please upload PNG files", "", [], None
    else:  # Batch Processing
        if zip_file:
            progress, gallery, zip_download = process_batch_zip(zip_file, fps, loop_count)
            return None, None, "", progress, gallery, zip_download
        else:
            return None, None, "", "Please upload a ZIP file", [], None

def create_gradio_interface():
    with gr.Blocks(title="PNG to GIF Converter") as demo:
        gr.Markdown("# PNG to GIF Converter")
        gr.Markdown("Create animated GIFs from PNG sequences")
        
        # Mode selector
        mode_selector = gr.Radio(
            choices=["Single Sequence", "Batch Processing"], 
            value="Single Sequence",
            label="Processing Mode"
        )
        gr.Markdown("**Single**: Upload individual PNG files | **Batch**: Upload ZIP with subdirectories")
        
        with gr.Row():
            with gr.Column():
                # Single mode components
                with gr.Group(visible=True) as single_group:
                    files_input = gr.File(
                        label="Upload PNG Files",
                        file_count="multiple",
                        file_types=[".png"]
                    )
                
                # Batch mode components  
                with gr.Group(visible=False) as batch_group:
                    gr.Markdown("ZIP should contain subdirectories, each with PNG files")
                    zip_input = gr.File(
                        label="Upload ZIP File",
                        file_count="single", 
                        file_types=[".zip"]
                    )
                
                # Common controls
                with gr.Row():
                    fps_input = gr.Slider(
                        minimum=1,
                        maximum=60,
                        value=12,
                        step=1,
                        label="Frames Per Second (FPS)"
                    )
                    
                    loop_input = gr.Slider(
                        minimum=0,
                        maximum=10,
                        value=0,
                        step=1,
                        label="Loop Count (0 = infinite)"
                    )
                
                convert_btn = gr.Button("Convert to GIF", variant="primary")
            
            with gr.Column():
                # Single mode outputs
                with gr.Group(visible=True) as single_output:
                    gif_preview = gr.Image(label="GIF Preview", type="filepath")
                    output_file = gr.File(label="Download GIF")
                
                # Batch mode outputs
                with gr.Group(visible=False) as batch_output:
                    batch_progress = gr.Textbox(label="Progress", interactive=False)
                    gif_gallery = gr.Gallery(label="Generated GIFs", columns=3, rows=2)
                    batch_download = gr.File(label="Download All GIFs (ZIP)")
                
                status_text = gr.Textbox(label="Status", interactive=False)
        
        # Mode switching logic
        mode_selector.change(
            fn=toggle_mode,
            inputs=[mode_selector],
            outputs=[single_group, batch_group, single_output, batch_output]
        )
        
        # Routing function to handle both modes
        convert_btn.click(
            fn=process_conversion,
            inputs=[mode_selector, files_input, zip_input, fps_input, loop_input],
            outputs=[gif_preview, output_file, status_text, batch_progress, gif_gallery, batch_download]
        )
    
    return demo

def main():
    # Check if running in a Hugging Face Space
    if os.getenv("SYSTEM") == "spaces":
        demo = create_gradio_interface()
        demo.launch()
        return

    ap = argparse.ArgumentParser()
    ap.add_argument("--src", default=".", help="Folder with PNG frames or parent directory for batch mode")
    ap.add_argument("--out", default="animation.gif", help="Output GIF file")
    ap.add_argument("--fps", type=int, default=12, help="Frames per second")
    ap.add_argument("--loop", type=int, default=0, help="Loop count (0 = infinite)")
    ap.add_argument("--web", action="store_true", help="Launch web interface")
    ap.add_argument("--batch", action="store_true", help="Batch process subdirectories")
    args = ap.parse_args()

    if args.web:
        demo = create_gradio_interface()
        demo.launch(share=True)
    elif args.batch:
        process_batch_directories(args.src, args.fps, args.loop)
    else:
        pngs = gather_frames(args.src)
        if not pngs:
            print(f"❌ No PNG files found in {args.src}")
            return
        build_gif(pngs, args.out, args.fps, args.loop)

if __name__ == "__main__":
    main()