import os
import glob
from PIL import Image, ImageDraw

def crop_to_16_9(image):
    """Crop image to 16:9 aspect ratio, centering the crop."""
    width, height = image.size
    target_ratio = 16 / 9
    current_ratio = width / height
    
    if current_ratio > target_ratio:
        # Image is too wide, crop width
        new_width = int(height * target_ratio)
        left = (width - new_width) // 2
        right = left + new_width
        cropped = image.crop((left, 0, right, height))
    else:
        # Image is too tall, crop height
        new_height = int(width / target_ratio)
        top = (height - new_height) // 2
        bottom = top + new_height
        cropped = image.crop((0, top, width, bottom))
    
    return cropped

def resize_to_target(image, target_width=2560, target_height=1440):
    """Resize image to target dimensions."""
    return image.resize((target_width, target_height), Image.Resampling.LANCZOS)

def overlay_logo(image, logo_path):
    """Overlay logo on the upper-right corner of the image."""
    if not os.path.exists(logo_path):
        print(f"Warning: Logo file {logo_path} not found. Skipping logo overlay.")
        return image
    
    try:
        logo = Image.open(logo_path).convert("RGBA")
        
        # Resize logo to be exactly 480 pixels tall while maintaining aspect ratio
        logo_height = 480
        logo_ratio = logo.size[0] / logo.size[1]  # width / height
        logo_width = int(logo_height * logo_ratio)
        logo = logo.resize((logo_width, logo_height), Image.Resampling.LANCZOS)
        
        # Get image dimensions
        img_width, img_height = image.size
        
        # Position logo in exact upper-right corner with no padding
        x = img_width - logo_width
        y = 0
        
        # Create a copy of the image to work with
        result = image.copy().convert("RGBA")
        
        # Paste logo with transparency
        result.paste(logo, (x, y), logo)
        
        # Convert back to RGB
        return result.convert("RGB")
    
    except Exception as e:
        print(f"Error overlaying logo: {e}")
        return image

def process_images():
    """Process all .jpg and .png images in the input folder."""
    input_folder = "input"
    output_folder = "output"
    logo_folder = "logo"
    
    # Create output folder if it doesn't exist
    os.makedirs(output_folder, exist_ok=True)
    
    # Find logo file in logo directory
    logo_files = glob.glob(os.path.join(logo_folder, "*.png"))
    if not logo_files:
        # Fallback to cornerlogo.png in root directory
        if os.path.exists("cornerlogo.png"):
            logo_path = "cornerlogo.png"
            print("Using cornerlogo.png from root directory")
        else:
            logo_path = None
            print("Warning: No logo file found")
    else:
        logo_path = logo_files[0]  # Use first PNG found in logo directory
        print(f"Using logo: {logo_path}")
    
    # Find all image files in input folder
    image_extensions = ["*.jpg", "*.jpeg", "*.png", "*.JPG", "*.JPEG", "*.PNG"]
    image_files = []
    
    for extension in image_extensions:
        image_files.extend(glob.glob(os.path.join(input_folder, extension)))
    
    if not image_files:
        print(f"No image files found in {input_folder} folder.")
        return
    
    print(f"Found {len(image_files)} image(s) to process:")
    
    for image_path in image_files:
        try:
            print(f"Processing: {os.path.basename(image_path)}")
            
            # Open and process the image
            with Image.open(image_path) as img:
                # Convert to RGB if necessary
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                
                # Step 1: Crop to 16:9 aspect ratio
                cropped_img = crop_to_16_9(img)
                
                # Step 2: Resize to 2560x1440
                resized_img = resize_to_target(cropped_img)
                
                # Step 3: Overlay logo
                if logo_path:
                    logo_img = overlay_logo(resized_img, logo_path)
                else:
                    logo_img = resized_img
                
                # Step 4: Save as JPG in output folder
                base_name = os.path.splitext(os.path.basename(image_path))[0]
                output_path = os.path.join(output_folder, f"{base_name}_logo.jpg")
                
                logo_img.save(output_path, "JPEG", quality=95)
                print(f"Saved: {output_path}")
        
        except Exception as e:
            print(f"Error processing {image_path}: {e}")

def main():
    """Main function to process images."""
    print("TeamsFlipper - Image Processing Tool")
    print("=" * 40)
    process_images()
    print("Processing complete!")

if __name__ == '__main__':
    main()