#!/usr/bin/env python3
"""AI background removal for product photos using rembg.

Originals are never modified. The script writes transparent PNG copies and a
manifest. It is intentionally separate from the older Pillow edge-removal
script so samples can be compared before replacing generated cutouts.
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

from PIL import Image

from rembg import new_session, remove


SUPPORTED = {".jpg", ".jpeg", ".png", ".webp"}


def iter_images(source: Path):
    for path in sorted(source.rglob("*")):
        if path.is_file() and path.suffix.lower() in SUPPORTED:
            yield path


def load_images(args) -> list[Path]:
    if args.image:
        return [Path(path) if Path(path).is_absolute() else Path.cwd() / path for path in args.image]

    if args.list_file:
        base = Path.cwd()
        return [
            Path(line.strip()) if Path(line.strip()).is_absolute() else base / line.strip()
            for line in Path(args.list_file).read_text(encoding="utf-8").splitlines()
            if line.strip()
        ]

    return list(iter_images(Path(args.source)))


def process_image(source: Path, output: Path, session, alpha_matting: bool) -> dict:
    if not source.exists() or source.suffix.lower() not in SUPPORTED:
        return {"source": str(source), "output": str(output), "skipped": "missing_or_unsupported"}

    with Image.open(source) as image:
        image = image.convert("RGBA")
        image.load()

    result = remove(
        image,
        session=session,
        alpha_matting=alpha_matting,
        alpha_matting_foreground_threshold=220,
        alpha_matting_background_threshold=15,
        alpha_matting_erode_size=3,
        post_process_mask=True,
    ).convert("RGBA")

    output.parent.mkdir(parents=True, exist_ok=True)
    result.save(output, "PNG", optimize=True)

    alpha = result.getchannel("A")
    bbox = alpha.getbbox()
    transparent = sum(1 for value in alpha.getdata() if value == 0)
    total = result.width * result.height

    return {
        "source": str(source),
        "output": str(output),
        "width": result.width,
        "height": result.height,
        "alpha_bbox": bbox,
        "transparent_ratio": round(transparent / total, 4) if total else 0,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source", default="storage/app/public/products")
    parser.add_argument("--output", default="storage/app/public/products-cutout-ai")
    parser.add_argument("--manifest", default="storage/app/product-cutouts-ai-manifest.json")
    parser.add_argument("--model-dir", default="storage/app/rembg-models")
    parser.add_argument("--model", default="isnet-general-use")
    parser.add_argument("--image", action="append", default=[])
    parser.add_argument("--list-file", default="")
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--skip-existing", action="store_true")
    parser.add_argument("--no-alpha-matting", action="store_true")
    args = parser.parse_args()

    model_dir = Path(args.model_dir)
    model_dir.mkdir(parents=True, exist_ok=True)
    os.environ.setdefault("U2NET_HOME", str(model_dir.resolve()))

    source_root = Path(args.source).resolve()
    output_root = Path(args.output)
    images = load_images(args)
    if args.limit:
        images = images[: args.limit]

    session = new_session(args.model)
    records = []

    for source in images:
        try:
            relative = source.resolve().relative_to(source_root)
        except ValueError:
            relative = Path(source.name)
        output = output_root / relative.with_suffix(".png")

        if args.skip_existing and output.exists():
            records.append({"source": str(source), "output": str(output), "skipped": "exists"})
            continue

        records.append(process_image(source, output, session, not args.no_alpha_matting))

    manifest = Path(args.manifest)
    manifest.parent.mkdir(parents=True, exist_ok=True)
    manifest.write_text(json.dumps({
        "model": args.model,
        "model_dir": str(model_dir),
        "source_root": str(source_root),
        "output_root": str(output_root),
        "count": len(records),
        "records": records,
    }, ensure_ascii=False, indent=2), encoding="utf-8")

    print(f"Processed {len(records)} images with {args.model}")
    print(f"Manifest: {manifest}")
    return 0


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