#!/usr/bin/env bash
# Create the isolated environment for the Polish number normalizer.
#
#   ./setup.sh                 venv + dependencies
#   ./setup.sh --model         also download both models (see below)
#   ./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"      # gated: times, dates, IDs
LLM_REPO="CYFRAGOVPL/PLLuM-4B-instruct-2512"            # open: grammar of plain integers
LLM_DIR="$HERE/models/PLLuM-4B-instruct-2512"
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
    "$VENV/bin/pip" install --quiet "huggingface_hub[cli]"

    # The inflection model first: it needs no account, so the install gets as
    # far as it can before anything can ask for credentials.
    if [ -f "$LLM_DIR/config.json" ]; then
        say "$LLM_REPO already present"
    else
        say "Downloading $LLM_REPO (8 GB, no account needed)"
        "$VENV/bin/hf" download "$LLM_REPO" --local-dir "$LLM_DIR"
    fi

    say "Downloading $MODEL_REPO"

    # 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 [ -f "$MODEL_DIR/config.json" ]; then
        say "$MODEL_REPO already present"
    elif [ -z "${HF_TOKEN:-}" ] && ! "$VENV/bin/hf" auth whoami >/dev/null 2>&1; then
        # Not fatal. Everything else is installed and useful; this one model
        # needs a human to accept a licence, which no installer can do.
        cat >&2 <<EOF

  [!] Skipped $MODEL_REPO — no Hugging Face credentials.
      It is gated (approval is automatic, but your account has to ask):

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

      Until then times, dates and identifiers fall back to num2words.

EOF
    else
        "$VENV/bin/hf" download "$MODEL_REPO" --local-dir "$MODEL_DIR"
    fi
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