#!/usr/bin/env python3
import sys
import os
from rembg import remove
from PIL import Image

"""
CLI usage:
  python3 rembg_remove_bg.py <input_image_path> <output_image_path>

Reads the input image, removes the background using rembg (U2Net), and
writes out a PNG with alpha channel to <output_image_path>.
Returns exit code 0 on success, non-zero on error.
"""

def main():
    if len(sys.argv) < 3:
        print("Usage: rembg_remove_bg.py <input_image> <output_image>", file=sys.stderr)
        return 2

    inp = sys.argv[1]
    out = sys.argv[2]

    try:
        # Ensure input exists
        if not os.path.isfile(inp):
            print(f"Input file not found: {inp}", file=sys.stderr)
            return 3

        # Create output directory if needed
        os.makedirs(os.path.dirname(out) or ".", exist_ok=True)

        # Open input image and remove background
        with Image.open(inp) as im:
            im = im.convert("RGBA")
            result = remove(im)
            # Always write PNG with alpha
            # Use optimize to reduce size, keep quality reasonable
            result.save(out, format="PNG", optimize=True)
        return 0
    except Exception as e:
        print(f"rembg error: {e}", file=sys.stderr)
        return 1


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