#!/usr/bin/env bash
set -euo pipefail

die() {
  echo "Error: $*" >&2
  exit 1
}

require_command() {
  command -v "$1" >/dev/null 2>&1 || die "missing required command: $1"
}

ensure_parent_dir() {
  mkdir -p "$(dirname "$1")"
}

usage() {
  cat <<'EOF'
Usage:
  remove-background.sh [--model NAME] [--install-only] <input> <output.png>

Examples:
  remove-background.sh assets/photo.jpg assets/photo.nobg.png
  remove-background.sh --model u2netp assets/photo.jpg assets/photo.fast.nobg.png
  remove-background.sh --install-only

Environment:
  REMBG_VENV_DIR   Path to the repo-local virtualenv. Default: <repo>/.venv-imagegen
  REMBG_MODEL_DIR  Path to the rembg model cache. Default: $HOME/.u2net
EOF
}

script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
repo_root="$(cd "$script_dir/../.." && pwd)"

venv_dir="${REMBG_VENV_DIR:-$repo_root/.venv-imagegen}"
model_dir="${REMBG_MODEL_DIR:-$HOME/.u2net}"
python_bin="$venv_dir/bin/python"
rembg_bin="$venv_dir/bin/rembg"

model_name="u2net"
install_only=0

while [ "$#" -gt 0 ]; do
  case "$1" in
    --model)
      shift
      [ "$#" -gt 0 ] || die "missing value for --model"
      model_name="$1"
      shift
      ;;
    --install-only)
      install_only=1
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    --)
      shift
      break
      ;;
    -*)
      die "unknown option: $1"
      ;;
    *)
      break
      ;;
  esac
done

if [ "$install_only" -eq 1 ]; then
  [ "$#" -eq 0 ] || die "--install-only does not take input/output arguments"
else
  [ "$#" -eq 2 ] || {
    usage
    exit 1
  }
fi

input_file="${1:-}"
output_file="${2:-}"

ensure_venv() {
  require_command python3

  if [ ! -x "$python_bin" ]; then
    python3 -m venv "$venv_dir"
  fi

  if ! "$python_bin" -m pip --version >/dev/null 2>&1; then
    "$python_bin" -m ensurepip --upgrade >/dev/null
  fi
}

ensure_rembg() {
  if ! "$python_bin" -c "import rembg" >/dev/null 2>&1; then
    "$python_bin" -m pip install 'rembg[cpu,cli]'
  fi

  [ -x "$rembg_bin" ] || die "rembg CLI was not installed into $venv_dir"
}

known_model_filename() {
  case "$1" in
    u2net)
      printf '%s\n' "u2net.onnx"
      ;;
    u2netp)
      printf '%s\n' "u2netp.onnx"
      ;;
    *)
      return 1
      ;;
  esac
}

known_model_url() {
  case "$1" in
    u2net)
      printf '%s\n' "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx"
      ;;
    u2netp)
      printf '%s\n' "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx"
      ;;
    *)
      return 1
      ;;
  esac
}

known_model_md5() {
  case "$1" in
    u2net)
      printf '%s\n' "60024c5c889badc19c04ad937298a77b"
      ;;
    u2netp)
      printf '%s\n' "8e83ca70e441ab06c318d82300c84806"
      ;;
    *)
      return 1
      ;;
  esac
}

prefetch_known_model() {
  local model="$1"
  local fname=""
  local url=""
  local expected_md5=""
  local target_path=""
  local tmp_path=""
  local current_md5=""

  if ! fname="$(known_model_filename "$model")"; then
    return 0
  fi

  url="$(known_model_url "$model")"
  expected_md5="$(known_model_md5 "$model")"
  target_path="$model_dir/$fname"
  tmp_path="$target_path.part"

  mkdir -p "$model_dir"

  if [ -f "$target_path" ] && command -v md5sum >/dev/null 2>&1; then
    current_md5="$(md5sum "$target_path" | awk '{print $1}')"
    if [ "$current_md5" = "$expected_md5" ]; then
      return 0
    fi
    rm -f "$target_path"
  elif [ -f "$target_path" ]; then
    return 0
  fi

  require_command curl

  rm -f "$tmp_path"
  curl -L --retry 5 --retry-all-errors -o "$tmp_path" "$url"

  if command -v md5sum >/dev/null 2>&1; then
    current_md5="$(md5sum "$tmp_path" | awk '{print $1}')"
    [ "$current_md5" = "$expected_md5" ] || {
      rm -f "$tmp_path"
      die "downloaded model checksum mismatch for $model"
    }
  fi

  mv "$tmp_path" "$target_path"
}

run_rembg() {
  ensure_parent_dir "$output_file"
  "$rembg_bin" i -m "$model_name" "$input_file" "$output_file"
}

ensure_venv
ensure_rembg
prefetch_known_model "$model_name"

if [ "$install_only" -eq 1 ]; then
  printf 'rembg ready in %s with model %s\n' "$venv_dir" "$model_name"
  exit 0
fi

[ -f "$input_file" ] || die "input file not found: $input_file"

run_rembg
printf '%s\n' "$output_file"
