#!/usr/bin/env python3
"""Losslessly recompress JPEG and PNG files.

JPEG goes through mozjpeg's lossless optimizer, which rebuilds the entropy
coding without touching the DCT coefficients, so there is no generational loss.
PNG goes through oxipng, which re-deflates the same pixel data.

Every file is verified by decoding both versions and comparing raw pixels. A
result is kept only if the pixels are identical AND it is smaller.

By default a new file is written next to each input (NAME.shrunk.jpg) or under
--output-dir. Overwriting needs --in-place plus --backup-dir or --no-backup.

    img-shrink photo.jpg
    img-shrink -r figures/ --output-dir smaller/
    img-shrink -r figures/ --in-place --backup-dir ~/img-originals
"""
import argparse
import csv
import hashlib
import io
import os
import shutil
import sys

SYNC_DIRS = ["~/Dropbox", "~/Google Drive", "~/GoogleDrive", "~/OneDrive",
             "~/Nextcloud", "~/ownCloud", "~/Sync", "~/iCloudDrive", "~/pCloudDrive"]
EXTS = (".jpg", ".jpeg", ".png")


def is_within(path, parent):
    try:
        return (os.path.commonpath([os.path.abspath(path), os.path.abspath(parent)])
                == os.path.abspath(parent))
    except ValueError:
        return False


def pixel_digest(data):
    """SHA-256 of the decoded pixels: the actual proof that nothing changed."""
    from PIL import Image
    Image.MAX_IMAGE_PIXELS = None
    with Image.open(io.BytesIO(data)) as im:
        im.load()
        # Hash only dimensions + RGBA pixels, not the mode: oxipng may
        # legitimately drop a fully opaque alpha channel (RGBA -> RGB) while
        # the rendered pixels stay identical.
        h = hashlib.sha256()
        h.update(f"{im.size}".encode())
        h.update(im.convert("RGBA").tobytes())
    return h.hexdigest()


def optimise(data, ext, png_level):
    if ext in (".jpg", ".jpeg"):
        import mozjpeg_lossless_optimization as mjo
        return mjo.optimize(data)
    import oxipng
    return oxipng.optimize_from_memory(data, level=png_level,
                                       strip=oxipng.StripChunks.safe())


def destination(path, args):
    if args.in_place:
        return path
    if args.output_dir:
        return os.path.join(args.output_dir, os.path.relpath(path, args.root))
    stem, ext = os.path.splitext(path)
    return stem + args.suffix + ext


def process(path, args):
    ext = os.path.splitext(path)[1].lower()
    try:
        with open(path, "rb") as fh:
            src = fh.read()
    except OSError as e:
        return (path, 0, 0, "SKIP", f"unreadable: {type(e).__name__}")
    size_in = len(src)
    if size_in < args.min_bytes:
        return (path, size_in, size_in, "SKIP", "below --min-bytes")
    try:
        before = pixel_digest(src)
    except Exception as e:
        return (path, size_in, size_in, "SKIP", f"undecodable: {type(e).__name__}")
    try:
        out = optimise(src, ext, args.png_level)
    except Exception as e:
        return (path, size_in, size_in, "SKIP", f"optimise failed: {type(e).__name__}")
    if not out or len(out) >= size_in:
        return (path, size_in, size_in, "SKIP", "not smaller")
    try:
        after = pixel_digest(out)
    except Exception as e:
        return (path, size_in, size_in, "FAIL", f"result undecodable: {type(e).__name__}")
    if before != after:
        return (path, size_in, size_in, "FAIL", "pixel mismatch")
    if args.dry_run:
        return (path, size_in, len(out), "WOULD-SHRINK", "pixels verified")
    st = os.stat(path)
    dst = destination(path, args)
    if args.in_place and args.backup_dir:
        bak = os.path.join(args.backup_dir, os.path.relpath(path, args.root))
        os.makedirs(os.path.dirname(bak), exist_ok=True)
        shutil.copy2(path, bak)
    os.makedirs(os.path.dirname(dst) or ".", exist_ok=True)
    tmp = dst + ".img-shrink-tmp"
    with open(tmp, "wb") as fh:
        fh.write(out)
    os.replace(tmp, dst)
    os.utime(dst, (st.st_atime, st.st_mtime))
    return (path, size_in, len(out), "SHRUNK", f"pixels verified -> {dst}")


def main(argv=None):
    ap = argparse.ArgumentParser(prog="img-shrink",
                                 description="Losslessly recompress JPEG/PNG files, "
                                             "verifying decoded pixels before keeping a result.")
    ap.add_argument("inputs", nargs="+", help="image files, or directories with -r")
    ap.add_argument("-r", "--recursive", action="store_true")
    ap.add_argument("--output-dir", help="write results here, mirroring the input tree")
    ap.add_argument("--suffix", default=".shrunk", help="suffix for default output names")
    ap.add_argument("--in-place", action="store_true",
                    help="replace inputs; requires --backup-dir or --no-backup")
    ap.add_argument("--backup-dir")
    ap.add_argument("--no-backup", action="store_true")
    ap.add_argument("--report", help="CSV report path")
    ap.add_argument("--jobs", type=int, default=1)
    ap.add_argument("--png-level", type=int, default=4, help="oxipng level 0-6")
    ap.add_argument("--min-bytes", type=int, default=8192)
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args(argv)

    if args.in_place and args.output_dir:
        ap.error("--in-place and --output-dir are mutually exclusive")
    if (args.backup_dir or args.no_backup) and not args.in_place:
        ap.error("--backup-dir / --no-backup only apply with --in-place")

    targets = []
    for item in map(os.path.abspath, args.inputs):
        if os.path.isdir(item):
            if not args.recursive:
                sys.exit(f"{item} is a directory; pass -r/--recursive")
            for dirpath, _d, files in os.walk(item):
                for fn in files:
                    stem = os.path.splitext(fn)[0]
                    if fn.lower().endswith(EXTS) and not stem.endswith(args.suffix):
                        targets.append(os.path.join(dirpath, fn))
        elif os.path.isfile(item):
            targets.append(item)
        else:
            sys.exit(f"not found: {item}")
    targets = sorted(set(targets))
    parents = [p if os.path.isdir(p) else os.path.dirname(p) for p in map(os.path.abspath, args.inputs)]
    args.root = os.path.commonpath(parents)
    if args.output_dir:
        args.output_dir = os.path.abspath(args.output_dir)

    if args.in_place and not args.dry_run:
        if args.backup_dir:
            args.backup_dir = os.path.abspath(os.path.realpath(os.path.expanduser(args.backup_dir)))
            for d in SYNC_DIRS:
                d = os.path.realpath(os.path.expanduser(d))
                if os.path.isdir(d) and is_within(args.backup_dir, d):
                    sys.exit(f"refusing: --backup-dir {args.backup_dir} is inside synced folder {d}")
            if is_within(args.backup_dir, args.root):
                sys.exit("refusing: --backup-dir is inside the tree being processed")
            os.makedirs(args.backup_dir, exist_ok=True)
        elif not args.no_backup:
            sys.exit("refusing: --in-place needs --backup-dir DIR or --no-backup")

    rep = w = None
    if args.report:
        rep = open(args.report, "w", newline="")
        w = csv.writer(rep)
        w.writerow(["path", "bytes_in", "bytes_out", "status", "note"])

    counts, saved = {}, 0

    def record(res):
        p, a, b, status, note = res
        counts[status] = counts.get(status, 0) + 1
        if w:
            w.writerow([os.path.relpath(p, args.root), a, b, status, note])
        if len(targets) <= 20 or status == "FAIL":
            print(f"  {status:12s} {os.path.relpath(p, args.root)}  {a:,} -> {b:,}  {note}")
        return a - b

    if args.jobs > 1 and targets:
        from concurrent.futures import ProcessPoolExecutor
        with ProcessPoolExecutor(max_workers=args.jobs) as ex:
            for res in ex.map(process, targets, [args] * len(targets)):
                saved += record(res)
    else:
        for p in targets:
            saved += record(process(p, args))
    if rep:
        rep.close()
    print(f"{'DRY RUN: ' if args.dry_run else ''}saved {saved/2**20:.2f} MB")
    for k, v in sorted(counts.items()):
        print(f"  {k:14s} {v}")
    return 1 if counts.get("FAIL") else 0


if __name__ == "__main__":
    sys.exit(main())
