import os
from pathlib import Path
from PIL import Image, ImageOps

# --- CONFIGURATION ---
INPUT_FOLDER = r"C:\Users\Dell\Desktop\gradio\weddingimages-hackathon\WeddingImagesHAck"   # Replace with your input folder path
OUTPUT_FOLDER = r"C:\Users\Dell\Desktop\gradio\weddingimages-hackathon-compressed"   # Replace with where you want to save them
MAX_DIMENSION = 2048                            # Maximum width or height in pixels
QUALITY = 85                                    # Image quality (1-95). 85 is highly optimized
# ---------------------

SUPPORTED_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.webp', '.tiff'}

def compress_image(file_path: Path, output_path: Path, max_dim: int, quality: int):
    try:
        with Image.open(file_path) as img:
            # 1. Correct image rotation from camera EXIF data automatically
            img = ImageOps.exif_transpose(img)
            
            # 2. Calculate new dimensions preserving the aspect ratio
            width, height = img.size
            if max(width, height) > max_dim:
                if width > height:
                    new_width = max_dim
                    new_height = int((max_dim / width) * height)
                else:
                    new_height = max_dim
                    new_width = int((max_dim / height) * width)
                
                img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
            
            # 3. Handle transparency channels if saving formats change
            ext = file_path.suffix.lower()
            if ext in {'.jpg', '.jpeg'}:
                # Force RGB conversion if image has an alpha channel (like transparent PNG to JPG conversion)
                if img.mode in ('RGBA', 'LA'):
                    img = img.convert('RGB')
                img.save(output_path, "JPEG", quality=quality, optimize=True)
            elif ext == '.webp':
                img.save(output_path, "WEBP", quality=quality, optimize=True)
            elif ext == '.png':
                # PNG is lossless, so we compress the file storage space directly
                img.save(output_path, "PNG", optimize=True)
            else:
                img.save(output_path, img.format, optimize=True)
            
            return True
    except Exception as e:
        print(f"Failed to process {file_path.name}: {e}")
        return False

def main():
    input_path = Path(INPUT_FOLDER)
    output_path = Path(OUTPUT_FOLDER)
    
    if not input_path.exists():
        print(f"Error: The input folder '{INPUT_FOLDER}' does not exist.")
        return
        
    output_path.mkdir(parents=True, exist_ok=True)
    
    files = [f for f in input_path.iterdir() if f.suffix.lower() in SUPPORTED_EXTENSIONS]
    total_files = len(files)
    
    if total_files == 0:
        print("No supported images found in the input folder.")
        return
        
    print(f"Starting compression of {total_files} images...")
    successful = 0
    
    for idx, file in enumerate(files, 1):
        out_file = output_path / file.name
        
        # Check original file size
        orig_size_mb = file.stat().st_size / (1024 * 1024)
        
        success = compress_image(file, out_file, MAX_DIMENSION, QUALITY)
        
        if success:
            successful += 1
            new_size_mb = out_file.stat().st_size / (1024 * 1024)
            reduction = ((orig_size_mb - new_size_mb) / orig_size_mb) * 100
            print(f"[{idx}/{total_files}] Compressed: {file.name} "
                  f"({orig_size_mb:.2f}MB -> {new_size_mb:.2f}MB | -{reduction:.1f}%)")
            
    print(f"\nCompression complete! Successfully processed {successful}/{total_files} images.")
    print(f"Your optimized files are located in: {OUTPUT_FOLDER}")

if __name__ == "__main__":
    main()