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

if [[ $# -ne 2 ]]; then
  echo "usage: bash install.sh ABSOLUTE_VENV_DIR ABSOLUTE_PUBLIC_RUNTIME_DIR" >&2
  exit 2
fi
VENV_DIR="$1"
PUBLIC_DIR="$2"
[[ "$VENV_DIR" = /* && "$PUBLIC_DIR" = /* ]] || {
  echo "Both paths must be absolute." >&2
  exit 2
}
[[ ! -e "$VENV_DIR" && ! -e "$PUBLIC_DIR" ]] || {
  echo "Use fresh venv and public-runtime directories." >&2
  exit 2
}

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
PUBLIC_COMMIT=15cc71d53cc2e6e18a68ee607d5fb6ba9a99e344
PUBLIC_ARCHIVE_SHA256=65e1e74b26989af16ba2232486d461d500249f69ebaa164eb210b69ed17cb1f6
WORK_DIR="$(dirname -- "$VENV_DIR")"
export PYTHONNOUSERSITE=1
unset PYTHONPATH
export PIP_CONFIG_FILE=/dev/null
export PIP_EXTRA_INDEX_URL=
export PIP_DISABLE_PIP_VERSION_CHECK=1
export PIP_CACHE_DIR="${PIP_CACHE_DIR:-$WORK_DIR/pip-cache}"
export TMPDIR="${TMPDIR:-$WORK_DIR/tmp}"
export XDG_CACHE_HOME="${XDG_CACHE_HOME:-$WORK_DIR/xdg-cache}"
export MAX_JOBS="${MAX_JOBS:-4}"
mkdir -p "$PIP_CACHE_DIR" "$TMPDIR" "$XDG_CACHE_HOME" "$(dirname -- "$PUBLIC_DIR")"

echo "STAGE create isolated Python 3.12 environment"
python3.12 -m venv "$VENV_DIR"
PYTHON="$VENV_DIR/bin/python"

echo "STAGE install PyTorch 2.5.0 CUDA 12.1"
timeout 1800s "$PYTHON" -m pip install --timeout 60 --retries 3 \
  --index-url https://download.pytorch.org/whl/cu121 \
  --extra-index-url https://pypi.org/simple \
  'torch==2.5.0+cu121' 'torchvision==0.20.0+cu121'

echo "STAGE install pinned inference dependencies"
timeout 1800s "$PYTHON" -m pip install --timeout 60 --retries 3 \
  --index-url https://pypi.org/simple \
  -r "$SCRIPT_DIR/requirements-inference.txt"

echo "STAGE fetch fixed public DanceGRPO source"
ARCHIVE="$WORK_DIR/DanceGRPO-$PUBLIC_COMMIT.tar.gz"
timeout 600s curl --fail --location --retry 3 --retry-delay 2 \
  "https://github.com/XueZeyue/DanceGRPO/archive/$PUBLIC_COMMIT.tar.gz" \
  --output "$ARCHIVE"
printf '%s  %s\n' "$PUBLIC_ARCHIVE_SHA256" "$ARCHIVE" | sha256sum --check --status
mkdir "$PUBLIC_DIR"
tar --extract --gzip --file "$ARCHIVE" --directory "$PUBLIC_DIR" --strip-components=1 \
  --no-same-owner --no-same-permissions
test -f "$PUBLIC_DIR/fastvideo/models/hunyuan_hf/pipeline_hunyuan.py"
timeout 600s "$PYTHON" -m pip install --timeout 60 --retries 3 \
  --index-url https://pypi.org/simple --no-deps --editable "$PUBLIC_DIR"

echo "STAGE build fixed flash-attn and install required public loader import"
timeout 3600s "$PYTHON" -m pip install --timeout 60 --retries 3 \
  --index-url https://pypi.org/simple --no-build-isolation 'flash-attn==2.7.0.post2'
timeout 600s "$PYTHON" -m pip install --timeout 60 --retries 3 \
  --index-url https://pypi.org/simple --no-deps 'liger-kernel==0.4.1'

echo "STAGE verify inference imports and installed versions"
"$PYTHON" - <<'PY'
import av, diffusers, fastvideo, huggingface_hub, torch, transformers
from fastvideo.models.hunyuan_hf.pipeline_hunyuan import HunyuanVideoPipeline
from fastvideo.utils.load import load_transformer
assert torch.__version__ == '2.5.0+cu121', torch.__version__
assert torch.version.cuda == '12.1', torch.version.cuda
assert diffusers.__version__ == '0.32.0', diffusers.__version__
assert transformers.__version__ == '4.46.1', transformers.__version__
print('verified:', torch.__version__, diffusers.__version__, transformers.__version__,
      av.__version__, huggingface_hub.__version__, flush=True)
PY
echo "INSTALL COMPLETE: $VENV_DIR"