#!/usr/bin/env python3
"""
Generate EXIF Orientation test images with manipulated pixel data.

All output images should look identical (upright) when EXIF is supported,
but will appear rotated/flipped when EXIF is ignored.

Usage: python generate_orientation_tests.py <input_image> <output_directory> --manipulate
"""

import os
import sys
import piexif
from PIL import Image

def ensure_dir(path):
    os.makedirs(path, exist_ok=True)

def transform_pixels_for_orientation(img, orientation):
    """
    Transform pixel data so that the visual result is always upright.
    The EXIF tag tells the viewer how to interpret the pixels.
    
    We apply the INVERSE transformation of what the EXIF tag would apply.
    For mirrors: they are their own inverse, but order matters with rotation.
    """
    if orientation == 1:
        return img.copy()  # No transformation needed
    
    elif orientation == 2:
        # EXIF: Mirror horizontal
        return img.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
    
    elif orientation == 3:
        # EXIF: Rotate 180°
        return img.rotate(180, expand=True)
    
    elif orientation == 4:
        # EXIF: Mirror vertical
        return img.transpose(Image.Transpose.FLIP_TOP_BOTTOM)
    
    elif orientation == 5:
        # EXIF: Rotate 90° CW then Mirror horizontal
        # Inverse: Mirror horizontal then Rotate 270° CCW (90° CW in PIL)
        temp = img.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
        return temp.rotate(90, expand=True)  # 90 CCW = 270 CW inverse
    
    elif orientation == 6:
        # EXIF: Rotate 90° CW
        # Inverse: Rotate 270° CCW (90° CCW in PIL)
        return img.rotate(90, expand=True)
    
    elif orientation == 7:
        # EXIF: Rotate 270° CW then Mirror horizontal
        # Inverse: Mirror horizontal then Rotate 90° CCW (270° CCW in PIL)
        temp = img.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
        return temp.rotate(270, expand=True)  # 270 CCW = 90 CW inverse
    
    elif orientation == 8:
        # EXIF: Rotate 270° CW
        # Inverse: Rotate 90° CCW (270° CCW in PIL)
        return img.rotate(270, expand=True)
    
    else:
        # Orientation 9 is typically invalid, treat as 1
        return img.copy()

def set_orientation_tag(image_path, output_dir, target_orientation, manipulate_pixels=False):
    """
    Loads an image, optionally manipulates pixels, sets EXIF tag, saves.
    """
    img = Image.open(image_path)
    
    # Convert to RGB for JPEG compatibility
    if img.mode in ('RGBA', 'P', 'LA', 'CMYK'):
        img = img.convert('RGB')
    
    # Optionally manipulate pixel data
    if manipulate_pixels:
        img = transform_pixels_for_orientation(img, target_orientation)
    
    # Build EXIF structure
    exif_dict = {"0th": {}, "Exif": {}, "GPS": {}, "1st": {}, "thumbnail": None}
    exif_dict["0th"][piexif.ImageIFD.Orientation] = target_orientation
    exif_bytes = piexif.dump(exif_dict)
    
    # Generate output filename
    base_name = os.path.basename(image_path)
    name, ext = os.path.splitext(base_name)
    if manipulate_pixels:
        output_filename = os.path.join(output_dir, f"{name}_Ori{target_orientation}_Manp.jpg")
    else:
        output_filename = os.path.join(output_dir, f"{name}_Ori{target_orientation}.jpg")
    
    # Save as JPEG with EXIF
    try:
        img.save(output_filename, 'JPEG', exif=exif_bytes, quality=95)
        mode = "pixels manipulated" if manipulate_pixels else "tag only"
        print(f"Generated: {os.path.basename(output_filename)} (Tag: {target_orientation}) [{mode}]")
    except Exception as e:
        print(f"Error saving {output_filename}: {e}")

def main():
    import argparse
    
    parser = argparse.ArgumentParser(description='Generate EXIF orientation test images')
    parser.add_argument('input_image', help='Input image path')
    parser.add_argument('output_dir', help='Output directory')
    parser.add_argument('--manipulate', action='store_true', 
                        help='Manipulate pixel data according to orientation tag')
    parser.add_argument('--tags-only', action='store_true',
                        help='Only change EXIF tag, do not manipulate pixels (default)')
    
    args = parser.parse_args()
    
    if not os.path.exists(args.input_image):
        print(f"Error: Input file '{args.input_image}' not found.")
        sys.exit(1)
    
    ensure_dir(args.output_dir)
    
    print(f"Processing: {args.input_image}")
    print(f"Output Directory: {args.output_dir}\n")
    
    # Determine which mode to use
    if args.manipulate:
        mode = "MANIPULATE_PIXELS"
    elif args.tags_only:
        mode = "TAGS_ONLY"
    else:
        mode = "TAGS_ONLY"  # Default
    
    print(f"Mode: {mode}\n")
    print("Generating EXIF Orientation Tags (1-9)...\n")
    
    for orient in range(1, 10):
        try:
            set_orientation_tag(
                args.input_image, 
                args.output_dir, 
                orient,
                manipulate_pixels=args.manipulate
            )
        except Exception as e:
            print(f"Failed to process Orientation {orient}: {e}")
    
    print("\n" + "="*60)
    print("TESTING GUIDE:")
    print("="*60)
    if args.manipulate:
        print("✅ PIXELS MANIPULATED:")
        print("   - All images should look UPRIGHT (same visual)")
        print("   - EXIF tag tells viewer how to display")
        print("   - Test: Open in different apps - some may show rotated/flipped!")
    else:
        print("⚠️  TAGS ONLY (original pixels):")
        print("   - Tag indicates rotation, pixels stay upright")
        print("   - EXIF-supported apps will rotate them")
        print("   - Non-EXIF apps will show them upright (ignoring tag)")
    print("="*60)

if __name__ == "__main__":
    main()