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

PROJECT_ROOT=${PROJECT_ROOT:-/data/vla_sft}
CODE_ROOT=${CODE_ROOT:-$PROJECT_ROOT/baselines}
VENV_ROOT=${VENV_ROOT:-$PROJECT_ROOT/envs}
PIP_INDEX_URL=${PIP_INDEX_URL-https://pypi.tuna.tsinghua.edu.cn/simple}
PIP_TRUSTED_HOST=${PIP_TRUSTED_HOST-pypi.tuna.tsinghua.edu.cn}
PIP_FALLBACK_INDEX_URL=${PIP_FALLBACK_INDEX_URL-https://pypi.org/simple}
PIP_DISABLE_FALLBACK=${PIP_DISABLE_FALLBACK:-0}
TORCH_INDEX_URL=${TORCH_INDEX_URL-}
TORCH_FALLBACK_INDEX_URL=${TORCH_FALLBACK_INDEX_URL-}
PYTHON_BIN=${PYTHON_BIN:-python3}
DREAMTACVLA_PYTHON_BIN=${DREAMTACVLA_PYTHON_BIN:-}
USE_NETWORK_TURBO=${USE_NETWORK_TURBO:-auto}
DRY_RUN=0
PROFILE=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --profile|--baseline)
      PROFILE="$2"; shift 2 ;;
    --code-root)
      CODE_ROOT="$2"; shift 2 ;;
    --venv-root)
      VENV_ROOT="$2"; shift 2 ;;
    --dry-run)
      DRY_RUN=1; shift ;;
    *)
      echo "unknown arg: $1" >&2; exit 2 ;;
  esac
done

if [[ -z "$PROFILE" ]]; then
  echo "usage: $0 --profile dreamzero|motus|dreamtacvla [--dry-run]" >&2
  exit 2
fi

echo "[env] PIP_INDEX_URL=${PIP_INDEX_URL:-<default pip>}"
echo "[env] PIP_FALLBACK_INDEX_URL=${PIP_FALLBACK_INDEX_URL:-<disabled>}"
echo "[env] TORCH_INDEX_URL=${TORCH_INDEX_URL:-<pip index>}"
echo "[env] TORCH_FALLBACK_INDEX_URL=${TORCH_FALLBACK_INDEX_URL:-<pip fallback/default>}"
echo "[env] PYTHON_BIN=$PYTHON_BIN"

if [[ "$USE_NETWORK_TURBO" != "0" && -f /etc/network_turbo ]]; then
  # AutoDL helper for academic GitHub/HuggingFace/PyTorch access. Keep the
  # proxy value injected by the helper; do not hard-code 127.0.0.1:12798.
  # shellcheck disable=SC1091
  source /etc/network_turbo >/dev/null 2>&1 || true
fi

run() {
  echo "+ $*"
  if [[ "$DRY_RUN" -eq 0 ]]; then
    "$@"
  fi
}

pip_install() {
  local args=("$@")
  if [[ -n "$PIP_INDEX_URL" ]]; then
    if [[ -n "$PIP_TRUSTED_HOST" ]]; then
      run python -m pip install -i "$PIP_INDEX_URL" --trusted-host "$PIP_TRUSTED_HOST" "${args[@]}" && return 0
    else
      run python -m pip install -i "$PIP_INDEX_URL" "${args[@]}" && return 0
    fi
    if [[ "$DRY_RUN" -eq 0 && "$PIP_DISABLE_FALLBACK" -ne 1 && -n "$PIP_FALLBACK_INDEX_URL" ]]; then
      echo "[pip fallback] primary index failed: $PIP_INDEX_URL" >&2
      echo "[pip fallback] retry with: $PIP_FALLBACK_INDEX_URL" >&2
      python -m pip install -i "$PIP_FALLBACK_INDEX_URL" "${args[@]}"
      return 0
    fi
    return 1
  else
    run env -u PIP_INDEX_URL -u PIP_TRUSTED_HOST python -m pip install "${args[@]}"
  fi
}

torch_install() {
  local args=("$@")
  local primary="${TORCH_INDEX_URL:-}"
  local fallback="${TORCH_FALLBACK_INDEX_URL:-}"

  if [[ -n "$primary" ]]; then
    run python -m pip install --index-url "$primary" "${args[@]}" && return 0
    if [[ "$DRY_RUN" -eq 0 && "$PIP_DISABLE_FALLBACK" -ne 1 ]]; then
      if [[ -n "$fallback" ]]; then
        echo "[torch fallback] primary torch index failed: $primary" >&2
        echo "[torch fallback] retry with: $fallback" >&2
        python -m pip install --index-url "$fallback" "${args[@]}"
        return 0
      fi
      echo "[torch fallback] primary torch index failed: $primary" >&2
      echo "[torch fallback] retry with default pip index" >&2
      pip_install "${args[@]}"
      return 0
    fi
    return 1
  fi

  pip_install "${args[@]}"
}

select_venv_python() {
  VENV_CREATE_PYTHON="$PYTHON_BIN"

  if [[ "$PROFILE" != "dreamtacvla" ]]; then
    return 0
  fi

  if [[ -n "$DREAMTACVLA_PYTHON_BIN" ]]; then
    VENV_CREATE_PYTHON="$DREAMTACVLA_PYTHON_BIN"
    return 0
  fi

  if command -v python3.11 >/dev/null 2>&1; then
    VENV_CREATE_PYTHON=$(command -v python3.11)
    return 0
  fi

  if command -v conda >/dev/null 2>&1; then
    local base_dir="$VENV_ROOT/.python311_base"
    if [[ "$DRY_RUN" -eq 1 ]]; then
      echo "+ conda create -y -p $base_dir python=3.11"
    elif [[ ! -x "$base_dir/bin/python" ]]; then
      conda create -y -p "$base_dir" python=3.11
    fi
    VENV_CREATE_PYTHON="$base_dir/bin/python"
    return 0
  fi

  VENV_CREATE_PYTHON="$PYTHON_BIN"
}

check_python_for_profile() {
  if [[ "$DRY_RUN" -eq 1 ]]; then
    echo "+ python - <<'PY'  # check Python version for $PROFILE"
    return 0
  fi

  if [[ "$PROFILE" == "dreamtacvla" ]]; then
    local py_version
    py_version=$(python - <<'PY'
import sys
print(f"{sys.version_info.major}.{sys.version_info.minor}")
PY
)
    if [[ "$py_version" != "3.11" ]]; then
      echo "DreamTacVLA requires Python 3.11 in this suite because mujoco==2.3.7 has no reliable Python $py_version wheel." >&2
      echo "Install python3.11, set DREAMTACVLA_PYTHON_BIN=/path/to/python3.11, or make conda available and recreate $VENV_DIR." >&2
      exit 4
    fi
  fi
}

prepare_resnet18_cache() {
  local cache_root url dest tmp hash
  cache_root="${TORCH_HOME:-$HOME/.cache/torch}/hub/checkpoints"
  url="https://download.pytorch.org/models/resnet18-f37072fd.pth"
  dest="$cache_root/resnet18-f37072fd.pth"
  tmp="$dest.tmp"

  if [[ "$DRY_RUN" -eq 1 ]]; then
    echo "+ mkdir -p $cache_root"
    echo "+ curl -L --fail --retry 5 -o $tmp $url"
    echo "+ sha256sum $tmp  # must start with f37072fd"
    return 0
  fi

  if [[ -s "$dest" ]]; then
    return 0
  fi

  mkdir -p "$cache_root"
  if command -v curl >/dev/null 2>&1; then
    curl -L -k --fail --connect-timeout 20 --retry 5 --retry-delay 2 -o "$tmp" "$url"
  else
    python - "$url" "$tmp" <<'PY'
import ssl
import sys
import urllib.request

url, dest = sys.argv[1], sys.argv[2]
ctx = ssl._create_unverified_context()
with urllib.request.urlopen(url, context=ctx, timeout=120) as response, open(dest, "wb") as out:
    out.write(response.read())
PY
  fi

  hash=$(sha256sum "$tmp" | awk '{print $1}')
  if [[ "$hash" != f37072fd* ]]; then
    echo "Downloaded ResNet18 checkpoint hash mismatch: $hash" >&2
    rm -f "$tmp"
    exit 6
  fi
  mv "$tmp" "$dest"
}

run mkdir -p "$VENV_ROOT"
VENV_DIR="$VENV_ROOT/$PROFILE"
select_venv_python

if [[ ! -d "$VENV_DIR" ]]; then
  run "$VENV_CREATE_PYTHON" -m venv "$VENV_DIR"
fi

# shellcheck disable=SC1091
if [[ "$DRY_RUN" -eq 0 ]]; then
  source "$VENV_DIR/bin/activate"
else
  echo "+ source $VENV_DIR/bin/activate"
fi

check_python_for_profile

pip_install --upgrade pip setuptools wheel

case "$PROFILE" in
  dreamzero)
    torch_install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0
    if [[ -f "$CODE_ROOT/dreamzero/requirements.txt" ]]; then
      pip_install -r "$CODE_ROOT/dreamzero/requirements.txt"
    fi
    pip_install deepspeed==0.19.1 accelerate==1.14.0 bitsandbytes==0.49.2 "numpy<2" mujoco==3.9.0 dm_tree imageio imageio-ffmpeg matplotlib albumentations ftfy peft
    ;;
  motus)
    torch_install torch==2.10.0 torchvision==0.25.0 torchaudio==2.10.0
    if [[ -f "$CODE_ROOT/motus/requirements.txt" ]]; then
      pip_install -r "$CODE_ROOT/motus/requirements.txt"
    fi
    pip_install deepspeed==0.19.2 accelerate==1.14.0 mujoco==3.9.0 "zarr>=3.1,<4"
    ;;
  dreamtacvla)
    torch_install torch==2.12.1 torchvision==0.27.1
    if [[ -f "$CODE_ROOT/dreamtacvla/requirements.txt" ]]; then
      pip_install -r "$CODE_ROOT/dreamtacvla/requirements.txt"
    fi
    pip_install transformers==5.12.1 open_clip_torch==3.3.0 timm==1.0.27 opencv-python==4.13.0.92 h5py==3.16.0 mujoco==2.3.7 imageio imageio-ffmpeg diffusers
    prepare_resnet18_cache
    ;;
  *)
    echo "unsupported profile: $PROFILE" >&2
    exit 3
    ;;
esac

run python - <<'PY'
import torch
print("torch", torch.__version__)
print("cuda", torch.version.cuda)
print("cuda_available", torch.cuda.is_available())
PY

echo "[env] done: $VENV_DIR"