#!/usr/bin/env python3
"""
Script to convert sets of images from pydep/test_run/ folders into PDFs.
Each folder contains three images: icp_aligned.png, pdf_img.png, and wld_img.png
"""

import os
import sys
from PIL import Image
from pathlib import Path
import glob
import json


def convert_images_to_pdf():
    """
    Convert all image sets in pydep/test_run/ folders to PDFs.
    Each folder's three images will be combined into a single PDF.
    """
    # Base directory containing the test_run folders
    base_dir = Path("pydep/test_run")
    
    if not base_dir.exists():
        print(f"Error: Directory {base_dir} does not exist!")
        return
    
    # Get all subdirectories (time_* folders)
    time_folders = [d for d in base_dir.iterdir() if d.is_dir() and d.name.startswith('time_')]
    
    if not time_folders:
        print(f"No time_* folders found in {base_dir}")
        return
    
    print(f"Found {len(time_folders)} folders to process...")
    
    # Process each folder
    for folder in sorted(time_folders):
        print(f"Processing folder: {folder.name}")
        
        # Expected image files in each folder
        image_files = [
            folder / "icp_aligned.png",
        ]
        
        # Check if all required images exist
        missing_files = [f for f in image_files if not f.exists()]
        if missing_files:
            print(f"  Warning: Missing files in {folder.name}: {[f.name for f in missing_files]}")
            continue
        
        # Open all images
        images = []
        for img_file in image_files:
            try:
                img = Image.open(img_file)
                # Convert to RGB if necessary (for PDF compatibility)
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                images.append(img)
                print(f"  Loaded: {img_file.name} ({img.size[0]}x{img.size[1]})")
            except Exception as e:
                print(f"  Error loading {img_file.name}: {e}")
                continue
        
        if not images:
            print(f"  No valid images found in {folder.name}")
            continue
        
        # Create PDF filename
        pdf_filename = f"outputs/{folder.name}_images.pdf"
        
        # Ensure outputs directory exists
        os.makedirs("outputs", exist_ok=True)
        
        # Save as PDF
        try:
            images[0].save(
                pdf_filename,
                "PDF",
                resolution=100.0,
                save_all=True,
                append_images=images[1:] if len(images) > 1 else []
            )
            print(f"  Created PDF: {pdf_filename}")
        except Exception as e:
            print(f"  Error creating PDF for {folder.name}: {e}")
        
        # Close images to free memory
        for img in images:
            img.close()
    
    print("Conversion completed!")

def create_combined_pdf():
    """
    Create a single PDF containing all images from all folders.
    Each folder's images will be on separate pages.
    """
    base_dir = Path("pydep/test_run")
    
    if not base_dir.exists():
        print(f"Error: Directory {base_dir} does not exist!")
        return
    
    time_folders = [d for d in base_dir.iterdir() if d.is_dir() and d.name.startswith('time_')]
    
    if not time_folders:
        print(f"No time_* folders found in {base_dir}")
        return
    
    print(f"Creating combined PDF from {len(time_folders)} folders...")
    
    all_images = []
    
    # Collect all images from all folders
    for folder in sorted(time_folders):
        print(f"Collecting images from: {folder.name}")
        
        image_files = [
            folder / "icp_aligned.png",
        ]
        
        for img_file in image_files:
            if img_file.exists():
                try:
                    img = Image.open(img_file)
                    if img.mode != 'RGB':
                        img = img.convert('RGB')
                    all_images.append(img)
                    print(f"  Added: {img_file.name}")
                except Exception as e:
                    print(f"  Error loading {img_file.name}: {e}")
    
    if not all_images:
        print("No images found to combine!")
        return
    
    # Create combined PDF
    pdf_filename = "outputs/all_images_combined.pdf"
    os.makedirs("outputs", exist_ok=True)
    
    try:
        all_images[0].save(
            pdf_filename,
            "PDF",
            resolution=100.0,
            save_all=True,
            append_images=all_images[1:] if len(all_images) > 1 else []
        )
        print(f"Created combined PDF: {pdf_filename} with {len(all_images)} pages")
    except Exception as e:
        print(f"Error creating combined PDF: {e}")
    
    # Close all images
    for img in all_images:
        img.close()

if __name__ == "__main__":
    print("Image to PDF Converter")
    print("=" * 50)
    
    # Check if PIL is available
    try:
        from PIL import Image
    except ImportError:
        print("Error: PIL (Pillow) is required. Install it with: pip install Pillow")
        sys.exit(1)
    
    # Create individual PDFs for each folder
    print("\n1. Creating individual PDFs for each folder...")
    convert_images_to_pdf()
    
    # Create combined PDF
    print("\n2. Creating combined PDF with all images...")
    create_combined_pdf()
    
    print("\nAll done!")