#!/usr/bin/env bash
# Create the isolated environment for the Polish number normalizer.
#
#   ./setup.sh                 venv + dependencies
#   ./setup.sh --model         also download the model (needs HF_TOKEN)
#   ./setup.sh --cpu           CPU-only torch, for a machine without an NVIDIA GPU
#
# Kept separate from the ComfyUI environment on purpose: this needs
# transformers >= 5, and the TTS install must not be disturbed.

set -euo pipefail

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VENV="$HERE/venv"
MODEL_DIR="$HERE/model"
MODEL_REPO="Folx/qwen3-0.6b-pl-text-normalization"
TORCH_INDEX="https://download.pytorch.org/whl/cu128"   # same wheels as ComfyUI

WANT_MODEL=0
CPU_ONLY=0
for arg in "$@"; do
    case "$arg" in
        --model) WANT_MODEL=1 ;;
        --cpu)   CPU_ONLY=1 ;;
        -h|--help) sed -n '2,10p' "${BASH_SOURCE[0]}" | sed 's/^# \?//'; exit 0 ;;
        *) echo "Unknown option: $arg" >&2; exit 2 ;;
    esac
done

say() { printf '\n\033[1;36m==> %s\033[0m\n' "$1"; }

say "Creating $VENV"
python3 -m venv "$VENV"
"$VENV/bin/pip" install --upgrade pip --quiet

say "Installing torch"
if [ "$CPU_ONLY" = 1 ]; then
    "$VENV/bin/pip" install torch --index-url https://download.pytorch.org/whl/cpu
else
    "$VENV/bin/pip" install torch --index-url "$TORCH_INDEX"
fi

say "Installing transformers, accelerate, num2words"
"$VENV/bin/pip" install "transformers>=5.0.0" "accelerate>=0.26.0" "num2words>=0.5.13"

if [ "$WANT_MODEL" = 1 ]; then
    say "Downloading $MODEL_REPO"
    "$VENV/bin/pip" install --quiet "huggingface_hub[cli]"

    # Either an HF_TOKEN in the environment or a stored `hf auth login` works;
    # the CLI picks up both on its own. Only say something when there is
    # neither, because the failure then is a confusing "access denied".
    if [ -z "${HF_TOKEN:-}" ] && ! "$VENV/bin/hf" auth whoami >/dev/null 2>&1; then
        cat >&2 <<EOF
No Hugging Face credentials. The repo is gated (approval is automatic, but it
has to be your account that asks):

  1. Open https://huggingface.co/$MODEL_REPO while signed in and accept the terms
  2. Either   $VENV/bin/hf auth login
     or       HF_TOKEN=hf_... $0 --model

EOF
        exit 1
    fi
    "$VENV/bin/hf" download "$MODEL_REPO" --local-dir "$MODEL_DIR"
fi

say "Checking the installation"
"$VENV/bin/python" - <<'PY'
import sys
try:
    import torch, transformers
    print(f"  torch {torch.__version__}  cuda_available={torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"  device: {torch.cuda.get_device_name(0)}")
    print(f"  transformers {transformers.__version__}")
    if int(transformers.__version__.split('.')[0]) < 5:
        print("  WARNING: transformers < 5 produces wrong output with this model",
              file=sys.stderr)
except ImportError as exc:
    print(f"  {exc}", file=sys.stderr); sys.exit(1)
PY

"$VENV/bin/python" "$HERE/normalizer.py" --status

if [ ! -f "$MODEL_DIR/config.json" ]; then
    cat <<EOF

The environment is ready, but the model is not downloaded yet, so the
normalizer will run on the degraded num2words fallback.

  1. Accept the terms at https://huggingface.co/$MODEL_REPO
  2. HF_TOKEN=hf_... $HERE/setup.sh --model

EOF
else
    say "Ready. Try:  $VENV/bin/python $HERE/normalizer.py 'Wygrał 1247 złotych.'"
fi