from PIL import Image
from pathlib import Path

# =====================
# SETTINGS
# =====================

# Folder that contains your source images:
# - Use Path("input_images") if you keep a dedicated folder
# - Use Path(".") if images are in the same folder as this script
INPUT_DIR = Path(".")


OUTPUT_FILE = "collage_5x15.png"

ROWS = 5
COLS = 15
REQUIRED = ROWS * COLS

# Each tile will be resized to this square size (uniform grid)
TILE_SIZE = 300  # px

# Spacing between tiles
H_SPACING = 30   # px
V_SPACING = 30   # px

# White background (RGBA)
BG_COLOR = (255, 255, 255, 255)

# "Square-ish" selection threshold:
# 1.00 = perfect square, 1.20 = quite strict, 1.35 = more permissive
ASPECT_MAX = 1.35

# Allowed extensions
EXTS = {".png", ".jpg", ".jpeg", ".webp"}


# =====================
# HELPERS
# =====================
def list_images(folder: Path):
    return sorted([p for p in folder.iterdir() if p.is_file() and p.suffix.lower() in EXTS])


def is_squareish(w: int, h: int, aspect_max: float) -> bool:
    long_side = max(w, h)
    short_side = min(w, h)
    if short_side == 0:
        return False
    aspect = long_side / short_side
    return aspect <= aspect_max


def fit_into_square_rgba(img: Image.Image, size: int, bg_color=(255, 255, 255, 255)) -> Image.Image:
    """
    Resize an image preserving aspect ratio and place it centered into a square canvas.
    """
    img = img.convert("RGBA")
    w, h = img.size

    # Scale to fit inside the square
    scale = min(size / w, size / h)
    new_w = max(1, int(w * scale))
    new_h = max(1, int(h * scale))
    img = img.resize((new_w, new_h), Image.LANCZOS)

    # Create square tile and paste centered
    tile = Image.new("RGBA", (size, size), bg_color)
    x = (size - new_w) // 2
    y = (size - new_h) // 2
    tile.paste(img, (x, y), img)
    return tile


# =====================
# MAIN
# =====================
if not INPUT_DIR.exists():
    raise SystemExit(f"ERROR: Input folder not found: {INPUT_DIR.resolve()}\n"
                     f"Create it or change INPUT_DIR to Path('.')")

files = list_images(INPUT_DIR)
print(f"Found {len(files)} image files in: {INPUT_DIR.resolve()}")

if len(files) == 0:
    raise SystemExit("ERROR: No images found. Check the folder and file extensions (.png/.jpg/.jpeg/.webp).")

# --- Filter square-ish images ---
squareish = []
for p in files:
    with Image.open(p) as im:
        w, h = im.size
    if is_squareish(w, h, ASPECT_MAX):
        squareish.append(p)

print(f"Square-ish (aspect <= {ASPECT_MAX}): {len(squareish)}")

if len(squareish) == 0:
    raise SystemExit(
        "ERROR: No square-ish images matched.\n"
        "Try increasing ASPECT_MAX (e.g. 1.50) or verify your images really have a square-ish canvas."
    )

# --- Ensure we have exactly REQUIRED tiles (loop/pattern if needed) ---
if len(squareish) < REQUIRED:
    print(f"Not enough square-ish images for {ROWS}x{COLS} ({REQUIRED}). Using loop/pattern to fill.")
    squareish = (squareish * (REQUIRED // len(squareish) + 1))[:REQUIRED]
else:
    squareish = squareish[:REQUIRED]

# --- Build tiles (uniform square thumbnails) ---
tiles = []
for p in squareish:
    img = Image.open(p)
    tile = fit_into_square_rgba(img, TILE_SIZE, BG_COLOR)
    tiles.append(tile)

# --- Create final canvas ---
canvas_w = COLS * TILE_SIZE + (COLS - 1) * H_SPACING
canvas_h = ROWS * TILE_SIZE + (ROWS - 1) * V_SPACING
canvas = Image.new("RGBA", (canvas_w, canvas_h), BG_COLOR)

# --- Paste tiles in a grid ---
idx = 0
for r in range(ROWS):
    for c in range(COLS):
        x = c * (TILE_SIZE + H_SPACING)
        y = r * (TILE_SIZE + V_SPACING)
        canvas.paste(tiles[idx], (x, y), tiles[idx])
        idx += 1

# --- Save output ---
canvas.save(OUTPUT_FILE)
print(f"✅ Collage created: {Path(OUTPUT_FILE).resolve()}")