#!/usr/bin/env python3
"""Compare background-removal backends on the live.f12.rocks inbox uploads.

Outputs land next to each other with explicit names so Fränky can pick a winner:

  <stem>__rembg-u2net__no-alpha.png
  <stem>__rembg-u2net__alpha.png
  <stem>__rembg-default__no-alpha.png
  <stem>__rembg-default__alpha.png
  <stem>__rembg-birefnet-general__no-alpha.png
  <stem>__rembg-birefnet-general__alpha.png
  <stem>__withoutbg-open-weights.png

`rembg-default` = rembg.remove() with no session (library default model).
`rembg-u2net`   = rembg.new_session("u2net") explicitly.
withoutbg embeds matting in its open-weights graph — one output only.
"""

from __future__ import annotations

import os
import sys
import time
import traceback
from pathlib import Path

os.environ.setdefault("TQDM_DISABLE", "1")

INPUT_DIR = Path(os.environ.get("COMPARE_INPUT", "/data/inputs"))
OUTPUT_DIR = Path(os.environ.get("COMPARE_OUTPUT", "/data/outputs"))
EXTS = {".jpg", ".jpeg", ".png", ".webp"}


def log(msg: str) -> None:
    print(msg, flush=True)


def run_rembg_named(data: bytes, label: str, model: str | None, alpha: bool, out: Path) -> None:
    from rembg import new_session, remove

    t0 = time.monotonic()
    kwargs: dict = {"alpha_matting": alpha}
    if model is None:
        result = remove(data, **kwargs)
        model_note = "default"
    else:
        session = new_session(model)
        result = remove(data, session=session, **kwargs)
        model_note = model
    out.write_bytes(result)
    log(f"  OK  {label} model={model_note} alpha={alpha} {time.monotonic()-t0:.1f}s -> {out.name} ({out.stat().st_size} bytes)")


def run_withoutbg(path: Path, out: Path) -> None:
    from withoutbg import WithoutBG

    t0 = time.monotonic()
    model = WithoutBG.open_weights()
    result = model.remove_background(str(path))
    result.save(out)
    log(f"  OK  withoutbg-open-weights {time.monotonic()-t0:.1f}s -> {out.name} ({out.stat().st_size} bytes)")


def process_one(path: Path) -> None:
    stem = path.stem
    log(f"=== {path.name} ===")
    data = path.read_bytes()
    jobs = [
        ("rembg-u2net", "u2net", False),
        ("rembg-u2net", "u2net", True),
        ("rembg-default", None, False),
        ("rembg-default", None, True),
        ("rembg-birefnet-general", "birefnet-general", False),
        ("rembg-birefnet-general", "birefnet-general", True),
    ]
    for label, model, alpha in jobs:
        tag = "alpha" if alpha else "no-alpha"
        out = OUTPUT_DIR / f"{stem}__{label}__{tag}.png"
        if out.exists() and out.stat().st_size > 0:
            log(f"  SKIP {out.name} (exists)")
            continue
        try:
            run_rembg_named(data, label, model, alpha, out)
        except Exception as exc:  # noqa: BLE001
            err = OUTPUT_DIR / f"{stem}__{label}__{tag}.ERROR.txt"
            err.write_text(f"{type(exc).__name__}: {exc}\n\n{traceback.format_exc()}", encoding="utf-8")
            log(f"  FAIL {label} alpha={alpha}: {type(exc).__name__}: {exc}")

    out_w = OUTPUT_DIR / f"{stem}__withoutbg-open-weights.png"
    if out_w.exists() and out_w.stat().st_size > 0:
        log(f"  SKIP {out_w.name} (exists)")
    else:
        try:
            run_withoutbg(path, out_w)
        except Exception as exc:  # noqa: BLE001
            err = OUTPUT_DIR / f"{stem}__withoutbg-open-weights.ERROR.txt"
            err.write_text(f"{type(exc).__name__}: {exc}\n\n{traceback.format_exc()}", encoding="utf-8")
            log(f"  FAIL withoutbg: {type(exc).__name__}: {exc}")


def main() -> int:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    inputs = sorted(p for p in INPUT_DIR.iterdir() if p.is_file() and p.suffix.lower() in EXTS)
    if not inputs:
        log(f"No images in {INPUT_DIR}")
        return 1
    log(f"Comparing {len(inputs)} image(s); outputs -> {OUTPUT_DIR}")
    for path in inputs:
        process_one(path)
    log("DONE")
    return 0


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