compressme.py

#!/usr/bin/env python3
"""
compressme.py — Compress every image inside a specified directory
while preserving visual quality as much as possible.
"""

from __future__ import annotations

import argparse
import logging
import os
import shutil
import sys
import time
import traceback
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Dict, List, Tuple

from src.backend.engine import (
    scan_images, create_output_path, compress_image, CompressResult,
)
from src.backend.hashing import detect_duplicates, sha256_hash
from src.backend.reporter import (
    print_summary, print_benchmark,
    write_csv_report, write_json_report,
    write_failed_images, write_duplicates_file,
)
from src.backend.utils import format_bytes
from src.backend.constants import (
    DEDUPE_OFF, DEDUPE_HASH, DEDUPE_PERCEPTUAL,
    FORMAT_KEEP, FORMAT_WEBP, FORMAT_AVIF,
)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("compressme")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Compress images in a directory while preserving visual quality.",
        add_help=True,
    )
    parser.add_argument("--dir", required=False, help="Target directory containing images.")
    parser.add_argument("--quality", type=int, default=85, help="Base quality (default: 85).")
    parser.add_argument("--format", choices=[FORMAT_KEEP, FORMAT_WEBP, FORMAT_AVIF], default=FORMAT_KEEP, help="Output format.")
    parser.add_argument("--dedupe", choices=[DEDUPE_OFF, DEDUPE_HASH, DEDUPE_PERCEPTUAL], default=DEDUPE_HASH, help="Duplicate detection mode.")
    parser.add_argument("--dry-run", action="store_true", help="Scan and estimate without writing files.")
    parser.add_argument("--benchmark", action="store_true", help="Show benchmark stats.")
    parser.add_argument("--workers", type=int, default=0, help="Worker processes (0 = auto).")
    return parser


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    if not args.dir:
        parser.print_help()
        print("\nError: --dir is required.")
        sys.exit(1)

    root_dir = Path(args.dir).resolve()
    if not root_dir.is_dir():
        print(f"Error: directory does not exist: {root_dir}")
        sys.exit(1)

    output_dir = root_dir / "compressed"
    output_dir.mkdir(parents=True, exist_ok=True)

    quality = max(1, min(100, args.quality))
    output_format = args.format
    dedupe_mode = args.dedupe
    dry_run = args.dry_run
    benchmark = args.benchmark
    workers = args.workers if args.workers > 0 else os.cpu_count() or 1

    print("Scanning...")
    images = scan_images(root_dir)
    images = [p for p in images if not p.parent.name.startswith("compressed")]

    print(f"Found {len(images)} images.")
    if not images:
        print("Nothing to do.")
        return

    if dry_run:
        total_orig = sum(p.stat().st_size for p in images)
        est_saved = int(total_orig * 0.7)
        est_compressed = total_orig - est_saved
        print(f"\nDry-run estimate for {len(images)} images:")
        print(f"  Original size:         {format_bytes(total_orig)}")
        print(f"  Estimated compressed:  {format_bytes(est_compressed)}")
        print(f"  Estimated savings:     {format_bytes(est_saved)}")
        print(f"  Estimated reduction:   ~70%")
        return

    hash_map: Dict[str, List[Path]] = {}
    dup_entries: List[Dict[str, str]] = []
    if dedupe_mode != DEDUPE_OFF:
        hash_map, _, dup_entries = detect_duplicates(images, dedupe_mode)
        if dup_entries:
            print(f"  Found {len(dup_entries)} duplicate(s).")
            write_duplicates_file(output_dir, dup_entries)

    work_items: List[Tuple[Path, Path]] = []
    skip_count = 0
    dup_count = 0
    compressed_hashes: Dict[str, Path] = {}

    for img_path in images:
        out_path = create_output_path(img_path, root_dir, output_dir)

        if out_path.exists():
            try:
                if out_path.stat().st_mtime >= img_path.stat().st_mtime:
                    log.info("Skipping %s — already optimized.", img_path.name)
                    skip_count += 1
                    continue
            except Exception:
                pass

        if dedupe_mode != DEDUPE_OFF:
            h = sha256_hash(img_path)
            if h in compressed_hashes:
                src_compressed = compressed_hashes[h]
                dest_ext = src_compressed.suffix
                dup_out = out_path.with_suffix(dest_ext)
                dup_out.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(src_compressed, dup_out)
                dup_count += 1
                log.info("Duplicate detected: %s \u2192 identical to %s", img_path.name, compressed_hashes[h].name)
                continue
            compressed_hashes[h] = out_path

        work_items.append((img_path, out_path))

    total_original = sum(p.stat().st_size for p in images)
    start_time = time.time()
    results: List[CompressResult] = []
    failures: List[Dict[str, str]] = []

    executor_class = ProcessPoolExecutor if workers > 1 else ThreadPoolExecutor
    with executor_class(max_workers=workers) as executor:
        futures = {}
        for img_path, out_path in work_items:
            future = executor.submit(compress_image, img_path, out_path, quality, output_format, False)
            futures[future] = (img_path, out_path)

        try:
            from tqdm import tqdm
            pbar = tqdm(as_completed(futures), total=len(futures), desc="Compressing", unit="img")
        except ImportError:
            pbar = as_completed(futures)

        for future in pbar:
            img_path, out_path = futures[future]
            try:
                result = future.result()
                results.append(result)
                if result["success"]:
                    saved = result["bytes_saved"]
                    if saved > 0 and result["original_size"] > 0:
                        pct = result["percent_saved"]
                        if hasattr(pbar, 'set_postfix_str'):
                            pbar.set_postfix_str(
                                f"{result['original_file']} ... "
                                f"Saved {format_bytes(result['original_size'])} -> "
                                f"{format_bytes(result['compressed_size'])} ({pct:.0f}%)"
                            )
                else:
                    failures.append({
                        "filename": img_path.name, "path": str(img_path),
                        "format": img_path.suffix.lower().lstrip("."),
                        "error": result.get("error", "Unknown"), "exception": "",
                    })
            except Exception as exc:
                failures.append({
                    "filename": img_path.name, "path": str(img_path),
                    "format": img_path.suffix.lower().lstrip("."),
                    "error": str(exc), "exception": traceback.format_exc(),
                })
                results.append({
                    "original_file": str(img_path), "compressed_file": "",
                    "original_size": img_path.stat().st_size, "compressed_size": 0,
                    "bytes_saved": 0, "percent_saved": 0.0,
                    "original_format": img_path.suffix.lower().lstrip("."),
                    "output_format": "", "width": 0, "height": 0,
                    "duplicate": "No", "duplicate_of": "",
                    "compression_method": "", "elapsed_time": 0.0,
                    "success": False, "error": str(exc),
                })

    total_time = time.time() - start_time
    total_compressed = sum(r["compressed_size"] for r in results)

    if failures:
        write_failed_images(output_dir, failures)

    all_results: List[CompressResult] = list(results)
    for de in dup_entries:
        dup_path = Path(de["duplicate"])
        all_results.append({
            "original_file": de["duplicate"], "compressed_file": de["original"],
            "original_size": dup_path.stat().st_size if dup_path.exists() else 0,
            "compressed_size": 0, "bytes_saved": 0, "percent_saved": 0.0,
            "original_format": dup_path.suffix.lower().lstrip(".") if dup_path.exists() else "",
            "output_format": "", "width": 0, "height": 0,
            "duplicate": "Yes", "duplicate_of": de["original"],
            "compression_method": "copy", "elapsed_time": 0.0, "success": True,
        })

    write_csv_report(output_dir, all_results)
    write_json_report(output_dir, all_results, total_time)

    print_summary(all_results, total_original, total_compressed, total_time,
                  dup_count + len(dup_entries), skip_count, len(failures), output_dir)

    if benchmark:
        print_benchmark(results, total_time)


if __name__ == "__main__":
    main()