#!/usr/bin/env python3
"""
Crop Gemini's watermark (the small star in the bottom-right corner) from PNG
images without quality loss.

The watermark sits in roughly the bottom-right 150px square of each image, so
this script trims 150px off both the right and bottom edges and re-saves the
PNG losslessly. 2048×2048 → 1898×1898 — still well above the 1024×1024 minimum
that Z-Image-Turbo (and most LoRA trainers) expect.

Usage:
  python3 crop_watermark.py <folder>                   # crop every PNG in <folder>, in place
  python3 crop_watermark.py <folder> --output <out>    # write cropped copies to <out> instead
  python3 crop_watermark.py <folder> --dry-run         # report what would happen, change nothing
"""

import argparse
import sys
from pathlib import Path

try:
    from PIL import Image
except ImportError:
    print("Pillow is required: pip3 install Pillow")
    sys.exit(1)

CROP_PX = 150  # pixels trimmed from the right and bottom edges


def crop_watermark(src: Path, dst: Path):
    """Crop the watermark area and save as a lossless PNG."""
    img = Image.open(src)
    w, h = img.size

    if w <= CROP_PX * 2 or h <= CROP_PX * 2:
        print(f"  SKIP {src.name}: too small ({w}x{h})")
        return False

    cropped = img.crop((0, 0, w - CROP_PX, h - CROP_PX))
    new_w, new_h = cropped.size

    # Save as a lossless PNG; preserve dpi / icc_profile metadata if present.
    info = img.info.copy()
    cropped.save(dst, "PNG", **{k: v for k, v in info.items() if k in ("dpi", "icc_profile")})

    print(f"  {src.name}: {w}x{h} → {new_w}x{new_h} → {dst.name}")
    return True


def main():
    parser = argparse.ArgumentParser(description="Crop Gemini's watermark from PNG images")
    parser.add_argument("folder", type=Path, help="Folder containing the PNG images")
    parser.add_argument(
        "--output",
        type=Path,
        default=None,
        help="Write cropped images to this folder (defaults to overwriting the originals)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Report what would be cropped, change nothing",
    )
    args = parser.parse_args()

    if not args.folder.is_dir():
        print(f"Folder not found: {args.folder}")
        sys.exit(1)

    images = sorted(args.folder.glob("*.png"))
    if not images:
        print(f"No PNG files found in {args.folder}")
        sys.exit(1)

    out_dir = args.output or args.folder
    if args.output:
        out_dir.mkdir(parents=True, exist_ok=True)

    print(f"Found {len(images)} PNG file(s)\n")

    if args.dry_run:
        for src in images:
            img = Image.open(src)
            w, h = img.size
            print(f"  {src.name}: {w}x{h} → {w - CROP_PX}x{h - CROP_PX}")
        print("\n--dry-run: nothing was changed")
        return

    count = 0
    for src in images:
        dst = out_dir / src.name
        if crop_watermark(src, dst):
            count += 1

    print(f"\nDone: cropped {count}/{len(images)} file(s)")


if __name__ == "__main__":
    main()