import os
from PIL import Image, ImageDraw, ImageFont

def create_mosaic(image_paths, output_path, columns, rows, target_width, padding=20, bg_color=(245, 245, 245)):
    images = []
    for path in image_paths:
        img = Image.open(path)
        # Calculate new height to maintain aspect ratio
        aspect_ratio = img.height / img.width
        target_height = int(target_width * aspect_ratio)
        img = img.resize((target_width, target_height), Image.Resampling.LANCZOS)
        images.append(img)
    
    if not images:
        return

    # All images should be same size now
    img_w, img_h = images[0].size
    
    mosaic_w = (img_w * columns) + (padding * (columns + 1))
    mosaic_h = (img_h * rows) + (padding * (rows + 1))
    
    mosaic = Image.new('RGB', (mosaic_w, mosaic_h), bg_color)
    
    for i, img in enumerate(images):
        if i >= columns * rows:
            break
        col = i % columns
        row = i // columns
        
        x = padding + col * (img_w + padding)
        y = padding + row * (img_h + padding)
        
        mosaic.paste(img, (x, y))
        
        # Add a subtle shadow/border
        draw = ImageDraw.Draw(mosaic)
        draw.rectangle([x-1, y-1, x+img_w, y+img_h], outline=(200, 200, 200), width=1)
        
    mosaic.save(output_path, quality=95)
    print(f"Saved {output_path}")

if __name__ == '__main__':
    src_dir = 'temp_screenshots'
    out_dir = 'docs/screenshots'
    os.makedirs(out_dir, exist_ok=True)
    
    # Get the 8 numbered screenshots
    files = sorted([f for f in os.listdir(src_dir) if f.startswith('0') and f.endswith('.png')])
    paths = [os.path.join(src_dir, f) for f in files]
    
    # Copy them to docs/screenshots
    import shutil
    for f in files:
        shutil.copy(os.path.join(src_dir, f), os.path.join(out_dir, f))
    
    # Create large mosaic (4x2 grid, each image 800px wide)
    create_mosaic(paths, os.path.join(out_dir, 'mosaic_large.png'), columns=4, rows=2, target_width=800)
    
    # Create small mosaic (4x2 grid, each image 250px wide)
    create_mosaic(paths, os.path.join(out_dir, 'mosaic_small.png'), columns=4, rows=2, target_width=250)