#!/usr/bin/env python3
"""Compress site images to <= 100 KB while preserving original colors."""

from __future__ import annotations

import argparse
import io
import shutil
from pathlib import Path

from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
SITE = ROOT / "public" / "images" / "site"
BACKUP = SITE / "_originals"
SKIP = {"logo.jpeg", "README.txt"}
SKIP_PATTERNS = (" copy", "-1.jpg")


def should_skip(path: Path) -> bool:
    if path.name in SKIP or path.name.startswith("_"):
        return True
    return any(p in path.name for p in SKIP_PATTERNS)


def encode_jpeg(img: Image.Image, quality: int) -> bytes:
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=quality, optimize=True, subsampling=2)
    return buf.getvalue()


def compress_to_target(img: Image.Image, max_bytes: int) -> tuple[Image.Image, int]:
    working = img.convert("RGB")
    scale = 1.0

    while scale >= 0.35:
        if scale < 1.0:
            w, h = working.size
            nw, nh = max(1, int(w * scale)), max(1, int(h * scale))
            candidate = working.resize((nw, nh), Image.Resampling.LANCZOS)
        else:
            candidate = working

        lo, hi = 20, 92
        best_q = lo
        best_data = encode_jpeg(candidate, lo)

        while lo <= hi:
            mid = (lo + hi) // 2
            data = encode_jpeg(candidate, mid)
            if len(data) <= max_bytes:
                best_q = mid
                best_data = data
                lo = mid + 1
            else:
                hi = mid - 1

        if len(best_data) <= max_bytes:
            return Image.open(io.BytesIO(best_data)).convert("RGB"), best_q

        scale *= 0.88

    data = encode_jpeg(candidate, 20)
    return Image.open(io.BytesIO(data)).convert("RGB"), 20


def optimize_file(path: Path, max_bytes: int, dry_run: bool) -> tuple[int, int]:
    before = path.stat().st_size
    if before <= max_bytes:
        return before, before

    with Image.open(path) as img:
        out, quality = compress_to_target(img, max_bytes)

    if dry_run:
        buf = io.BytesIO()
        out.save(buf, format="JPEG", quality=quality, optimize=True, subsampling=2)
        return before, len(buf.getvalue())

    out.save(path, format="JPEG", quality=quality, optimize=True, subsampling=2)
    return before, path.stat().st_size


def main() -> None:
    parser = argparse.ArgumentParser(description="Optimize site images to <= 100 KB")
    parser.add_argument("--max-kb", type=int, default=100)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    max_bytes = args.max_kb * 1024
    BACKUP.mkdir(exist_ok=True)

    files = sorted(SITE.glob("*.jpg")) + sorted(SITE.glob("*.jpeg"))
    changed = 0

    for path in files:
        if should_skip(path):
            continue

        backup = BACKUP / path.name
        if not backup.exists() and not args.dry_run:
            shutil.copy2(path, backup)

        before, after = optimize_file(path, max_bytes, args.dry_run)
        if after < before or before > max_bytes:
            changed += 1
            print(
                f"{'~' if args.dry_run else '✓'} {path.name}: "
                f"{before // 1024} KB → {after // 1024} KB"
            )
        else:
            print(f"○ {path.name}: {before // 1024} KB (ok)")

    print(f"\nDone — {changed} optimized, originals in {BACKUP}")


if __name__ == "__main__":
    main()