#!/usr/bin/env python3
"""One-command local background remover (rembg) from GenLovers.

Usage:
    python setup_remove_background.py photo.jpg [more.png ...]

First run: creates a private virtual environment, installs rembg (CPU) into it,
then cuts out every image you pass. Later runs skip the install. The AI model
downloads once, on the first cutout. Everything runs on this machine; nothing is
uploaded anywhere. Needs Python 3.10 or newer (tested on 3.13) and an internet connection for
the first run.
"""
import os
import subprocess
import sys
import venv
from pathlib import Path

ROOT = Path(os.environ.get("GENLOVERS_TOOLS_DIR", Path.home() / "genlovers-tools"))
ENV_DIR = ROOT / "remove-background"
BIN = ENV_DIR / ("Scripts" if os.name == "nt" else "bin")
PY = BIN / ("python.exe" if os.name == "nt" else "python")
MODEL = "u2net"
REMBG = BIN / ("rembg.exe" if os.name == "nt" else "rembg")


def say(msg):
    print(f"\n==> {msg}", flush=True)


def fail(msg):
    print(f"\nSETUP FAILED: {msg}", file=sys.stderr)
    sys.exit(1)


def ensure_installed():
    if REMBG.exists():
        say("rembg already installed.")
        return
    if sys.version_info < (3, 10):
        fail("Python 3.10 or newer is required. Install it from https://www.python.org/downloads/")
    say(f"Creating a private environment in {ENV_DIR}")
    venv.create(ENV_DIR, with_pip=True)
    say("Installing rembg (this can take a few minutes)...")
    result = subprocess.run([str(PY), "-m", "pip", "install", "--upgrade", "rembg[cpu,cli]"])
    if result.returncode != 0 or not REMBG.exists():
        fail(
            "pip could not install rembg. Try Python 3.13, the version this "
            "script was tested on, and run it again."
        )


def main():
    images = [Path(a) for a in sys.argv[1:]]
    ensure_installed()
    if not images:
        print(
            "\nSETUP COMPLETE\nNow cut out an image with:\n"
            "    python setup_remove_background.py photo.jpg"
        )
        return
    failed = 0
    for image in images:
        if not image.is_file():
            print(f"Skipping {image}: file not found.", file=sys.stderr)
            failed += 1
            continue
        out = image.with_name(f"{image.stem}-cutout.png")
        say(f"Removing the background from {image.name} (the model downloads on the first image)...")
        if subprocess.run([str(REMBG), "i", "-m", MODEL, str(image), str(out)]).returncode == 0:
            print(f"Saved {out}")
        else:
            failed += 1
    sys.exit(1 if failed else 0)


if __name__ == "__main__":
    main()