main.py

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Optional

from rich.console import Console
from rich.table import Table

from config import Config
from pipeline import ImagePipeline, PipelineResult
from utils import human_size, read_urls_from_file, setup_logging


console = Console()


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="image-processor-pro",
        description="Download, convert, and optimize images to JPEG.",
    )
    source = parser.add_mutually_exclusive_group(required=False)
    source.add_argument("--url", help="A single image URL to process.")
    source.add_argument("--file", type=Path, help="Path to a text file with one URL per line.")

    parser.add_argument("--config", type=Path, help="Optional JSON config file overriding defaults.")
    parser.add_argument("--output", type=Path, help="Override output root directory.")
    parser.add_argument("--logs", type=Path, help="Override logs directory.")
    parser.add_argument("--quality", type=int, help="JPEG quality (1-95). Default: 85.")
    parser.add_argument("--max-width", type=int, help="Resize so width does not exceed this value.")
    parser.add_argument("--max-height", type=int, help="Resize so height does not exceed this value.")
    parser.add_argument("--workers", type=int, help="Number of concurrent download workers.")
    parser.add_argument("--timeout", type=int, help="Per-request download timeout in seconds.")
    parser.add_argument("--retries", type=int, help="Maximum download retries per URL.")
    parser.add_argument("--keep-metadata", action="store_true", help="Preserve EXIF metadata if present.")
    parser.add_argument("--remove-watermark", action="store_true",
                        help="Inpaint over a fixed bottom-left watermark region.")
    parser.add_argument("--wm-engine", choices=["lama", "classical"],
                        help="Inpaint engine: 'lama' (AI, default, handles textured backgrounds) "
                             "or 'classical' (cv2, fast, no extra deps).")
    parser.add_argument("--wm-device", choices=["cpu", "mps"],
                        help="LaMa compute device (default cpu). 'mps' uses the Apple Silicon GPU.")
    parser.add_argument("--wm-x", type=float, help="Watermark region left edge as fraction of width (0-1).")
    parser.add_argument("--wm-y", type=float, help="Watermark region top edge as fraction of height (0-1).")
    parser.add_argument("--wm-w", type=float, help="Watermark region width as fraction of image width (0-1).")
    parser.add_argument("--wm-h", type=float, help="Watermark region height as fraction of image height (0-1).")
    parser.add_argument("--wm-radius", type=int, help="cv2.inpaint radius in pixels.")
    parser.add_argument("--wm-kernel", type=int,
                        help="Morphological kernel size for text detection (odd, default 9).")
    parser.add_argument("--wm-threshold", type=int,
                        help="Contrast threshold 0-255 (default 25). Lower = more aggressive.")
    parser.add_argument("--wm-min-area", type=int,
                        help="Minimum connected-component area kept (default 4).")
    parser.add_argument("--wm-max-area-frac", type=float,
                        help="Reject text blobs larger than this fraction of the search ROI (default 0.6).")
    parser.add_argument("--wm-feather", type=int,
                        help="Mask dilation in pixels to cover edges (default 1).")
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose logging.")
    return parser


def load_config(args: argparse.Namespace) -> Config:
    config = Config()

    if args.config:
        config = _apply_config_file(config, args.config)

    if args.output:
        config.output_root = args.output
    if args.logs:
        config.logs_root = args.logs
    if args.quality is not None:
        config.jpeg_quality = max(1, min(95, args.quality))
    if args.max_width is not None:
        config.max_width = args.max_width
    if args.max_height is not None:
        config.max_height = args.max_height
    if args.workers is not None:
        config.max_workers = max(1, args.workers)
    if args.timeout is not None:
        config.download_timeout = max(1, args.timeout)
    if args.retries is not None:
        config.max_retries = max(1, args.retries)
    if args.keep_metadata:
        config.strip_metadata = False
    if args.remove_watermark:
        config.remove_watermark = True
    if args.wm_engine is not None:
        config.watermark_engine = args.wm_engine
    if args.wm_device is not None:
        config.watermark_device = args.wm_device
    if args.wm_x is not None:
        config.watermark_x_frac = args.wm_x
    if args.wm_y is not None:
        config.watermark_y_frac = args.wm_y
    if args.wm_w is not None:
        config.watermark_w_frac = args.wm_w
    if args.wm_h is not None:
        config.watermark_h_frac = args.wm_h
    if args.wm_radius is not None:
        config.watermark_inpaint_radius = max(1, args.wm_radius)
    if args.wm_kernel is not None:
        config.watermark_text_kernel = max(3, args.wm_kernel)
    if args.wm_threshold is not None:
        config.watermark_text_threshold = max(1, min(255, args.wm_threshold))
    if args.wm_min_area is not None:
        config.watermark_min_component_area = max(1, args.wm_min_area)
    if args.wm_max_area_frac is not None:
        config.watermark_max_area_frac = args.wm_max_area_frac
    if args.wm_feather is not None:
        config.watermark_mask_feather = max(0, args.wm_feather)

    return config


def _apply_config_file(config: Config, path: Path) -> Config:
    with path.open("r", encoding="utf-8") as fh:
        data = json.load(fh)

    field_map = {
        "output_root": Path,
        "logs_root": Path,
        "jpeg_quality": int,
        "max_width": lambda v: int(v) if v is not None else None,
        "max_height": lambda v: int(v) if v is not None else None,
        "download_timeout": int,
        "max_retries": int,
        "retry_backoff": float,
        "max_workers": int,
        "chunk_size": int,
        "user_agent": str,
        "strip_metadata": bool,
        "remove_watermark": bool,
        "watermark_engine": str,
        "watermark_device": str,
        "watermark_x_frac": float,
        "watermark_y_frac": float,
        "watermark_w_frac": float,
        "watermark_h_frac": float,
        "watermark_inpaint_radius": int,
        "watermark_mask_feather": int,
        "watermark_text_kernel": int,
        "watermark_text_threshold": int,
        "watermark_min_component_area": int,
        "watermark_max_area_frac": float,
    }
    for key, caster in field_map.items():
        if key in data:
            setattr(config, key, caster(data[key]))

    if "background_color" in data:
        config.background_color = tuple(int(c) for c in data["background_color"])

    return config


def collect_urls(args: argparse.Namespace) -> list[str]:
    if args.url:
        return [args.url]
    if args.file:
        if not args.file.exists():
            console.print(f"[red]URL file not found:[/red] {args.file}")
            sys.exit(2)
        return read_urls_from_file(args.file)
    console.print("[red]Provide --url or --file.[/red]")
    sys.exit(2)


def print_summary(results: list[PipelineResult]) -> None:
    successes = [r for r in results if r.success]
    failures = [r for r in results if not r.success]

    table = Table(title="Image Processor Pro — Summary", show_lines=False)
    table.add_column("Metric", style="bold cyan")
    table.add_column("Value", style="white")
    table.add_row("Total", str(len(results)))
    table.add_row("Succeeded", f"[green]{len(successes)}[/green]")
    table.add_row("Failed", f"[red]{len(failures)}[/red]")

    if successes:
        total_bytes = sum(r.final_size_bytes for r in successes)
        avg_dl = sum(r.download_seconds for r in successes) / len(successes)
        avg_conv = sum(r.convert_seconds for r in successes) / len(successes)
        avg_opt = sum(r.optimize_seconds for r in successes) / len(successes)
        table.add_row("Total output size", human_size(total_bytes))
        table.add_row("Avg download time", f"{avg_dl:.2f}s")
        table.add_row("Avg convert time", f"{avg_conv:.2f}s")
        table.add_row("Avg optimize time", f"{avg_opt:.2f}s")

    console.print(table)

    if failures:
        fail_table = Table(title="Failures", show_lines=False)
        fail_table.add_column("URL", style="yellow", overflow="fold")
        fail_table.add_column("Error", style="red", overflow="fold")
        for r in failures:
            fail_table.add_row(r.url, r.error or "unknown")
        console.print(fail_table)


def main(argv: Optional[list[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    config = load_config(args)
    setup_logging(config.logs_root, verbose=args.verbose)

    urls = collect_urls(args)
    if not urls:
        console.print("[yellow]No URLs to process.[/yellow]")
        return 0

    console.print(f"[bold]Processing[/bold] {len(urls)} URL(s) with {config.max_workers} workers")

    with ImagePipeline(config) as pipeline:
        results = pipeline.process_many(urls)

    print_summary(results)
    return 0 if all(r.success for r in results) else 1


if __name__ == "__main__":
    raise SystemExit(main())