#!/usr/bin/env python3
"""Offline, model-independent PNG optimization. Originals are never overwritten.

Usage: python optimize-image.py image.png [another.png ...]
Only a smaller, separate *.optimized.png file is saved. Other formats, animated
or signed PNGs, and unsafe/invalid inputs remain unchanged. Python 3.9+ only;
no packages, network, credentials, image generation, or provider APIs required.
"""
import argparse
import json
import os
from pathlib import Path
import struct
import sys
import zlib

PNG = b"\x89PNG\r\n\x1a\n"
MAX_FILE_BYTES = 64 * 1024 * 1024
MAX_RAW_BYTES = 128 * 1024 * 1024
MAX_CHUNKS = 4096
# Metadata is copied byte-for-byte, not decoded/re-encoded or stripped.
KNOWN_CHUNKS = {
    b"IHDR", b"PLTE", b"IDAT", b"IEND", b"cHRM", b"gAMA", b"iCCP",
    b"sBIT", b"sRGB", b"bKGD", b"hIST", b"tRNS", b"pHYs", b"tIME",
    b"iTXt", b"tEXt", b"zTXt", b"eXIf", b"cICP", b"mDCV", b"cLLI",
}


def optimize_png(data):
    """Recompress the exact filtered scanlines; do not quantize or change filters."""
    if not data.startswith(PNG):
        return data, "unsupported_format"
    if len(data) > MAX_FILE_BYTES:
        return data, "size_limit"
    chunks = []
    pos = len(PNG)
    while pos < len(data):
        if len(chunks) >= MAX_CHUNKS:
            return data, "size_limit"
        if pos + 12 > len(data):
            return data, "invalid_png"
        size = struct.unpack_from(">I", data, pos)[0]
        end = pos + size + 12
        if end > len(data):
            return data, "invalid_png"
        kind = data[pos + 4:pos + 8]
        body = data[pos + 8:end - 4]
        if zlib.crc32(kind + body) != struct.unpack_from(">I", data, end - 4)[0]:
            return data, "invalid_png"
        if kind in (b"acTL", b"fcTL", b"fdAT"):
            return data, "animated_png"
        if kind in (b"caBX", b"dSIG"):
            return data, "signed_png"
        if kind not in KNOWN_CHUNKS:
            return data, "unsupported_png_chunk"
        chunks.append((kind, body, data[pos:end]))
        pos = end
        if kind == b"IEND":
            break
    kinds = [c[0] for c in chunks]
    if (pos != len(data) or not chunks or kinds[0] != b"IHDR"
            or kinds.count(b"IHDR") != 1 or kinds[-1] != b"IEND"
            or chunks[-1][1] or len(chunks[0][1]) != 13 or b"IDAT" not in kinds):
        return data, "invalid_png"
    first, last = kinds.index(b"IDAT"), len(kinds) - 1 - kinds[::-1].index(b"IDAT")
    if any(k != b"IDAT" for k in kinds[first:last + 1]):
        return data, "invalid_png"
    width, height, depth, color, compression, filtering, interlace = struct.unpack(">IIBBBBB", chunks[0][1])
    depths = {0: (1, 2, 4, 8, 16), 2: (8, 16), 3: (1, 2, 4, 8), 4: (8, 16), 6: (8, 16)}
    if (not width or not height or depth not in depths.get(color, ())
            or compression or filtering or interlace not in (0, 1)):
        return data, "invalid_png"
    if interlace:
        return data, "interlaced_png"
    channels = {0: 1, 2: 3, 3: 1, 4: 2, 6: 4}[color]
    stride = (width * channels * depth + 7) // 8 + 1
    expected = stride * height
    if expected > MAX_RAW_BYTES:
        return data, "size_limit"
    try:
        decoder = zlib.decompressobj()
        raw = decoder.decompress(b"".join(c[1] for c in chunks[first:last + 1]), expected + 1)
        if (len(raw) != expected or not decoder.eof or decoder.unused_data
                or decoder.unconsumed_tail or any(raw[i] > 4 for i in range(0, len(raw), stride))):
            return data, "invalid_png"
        compressed = zlib.compress(raw, 9)
    except zlib.error:
        return data, "invalid_png"
    payload = b"IDAT" + compressed
    idat = struct.pack(">I", len(compressed)) + payload + struct.pack(">I", zlib.crc32(payload))
    candidate = PNG + b"".join(c[2] for c in chunks[:first]) + idat + b"".join(c[2] for c in chunks[last + 1:])
    if len(candidate) >= len(data):
        return data, "already_compact"
    return candidate, "optimized"


def optimize_file(path):
    """Return the selected local download path; any skip retains the original."""
    result = {"original": str(path), "download": str(path)}
    with path.open("rb") as source:
        # Bound reads even if the input grows while it is being read.
        data = source.read(MAX_FILE_BYTES + 1)
    if len(data) > MAX_FILE_BYTES:
        return {**result, "status": "unchanged", "reason": "size_limit"}
    candidate, reason = optimize_png(data)
    result.update(original_bytes=len(data), selected_bytes=len(candidate), bytes_saved=len(data) - len(candidate))
    if reason != "optimized":
        return {**result, "status": "unchanged", "reason": reason}
    output = path.with_name(path.stem + ".optimized.png")
    # Exclusive creation protects originals, existing derivatives, and symlinks.
    try:
        target = os.fdopen(os.open(output, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "wb")
    except FileExistsError:
        return {**result, "status": "unchanged", "reason": "output_exists", "selected_bytes": len(data), "bytes_saved": 0}
    try:
        with target:
            target.write(candidate)
    except OSError:
        output.unlink(missing_ok=True)
        raise
    return {**result, "download": str(output), "status": "optimized"}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("images", nargs="+", type=Path)
    args = parser.parse_args()
    failed = False
    for path in args.images:
        try:
            result = optimize_file(path)
        except OSError as exc:
            # No credentials are accessed; no provider or inference retry occurs.
            result = {"original": str(path), "download": str(path), "status": "error", "reason": type(exc).__name__}
            failed = True
        print(json.dumps(result, ensure_ascii=False))
    return int(failed)


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