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

PROFILE=dreamtacvla
BENCH=manifeel
MODE=smoke
DRY_RUN=0
TRAIN=1
TMUX_SESSION=""

usage() {
  cat >&2 <<'EOF'
usage: scripts/deploy_verified.sh [--profile dreamtacvla] [--bench manifeel] [--smoke|--full|--prepare-only] [--dry-run] [--tmux SESSION]

Verified deployment entrypoint.

Currently verified end-to-end:
  dreamtacvla + manifeel

Default paths:
  PROJECT_ROOT=/root/autodl-tmp/vla_sft

Useful overrides:
  PROJECT_ROOT DATA_ROOT CKPT_ROOT CODE_ROOT VENV_ROOT OUTPUT_ROOT
  HF_USERNAME HF_TOKEN USE_NETWORK_TURBO
  ARIA2_CONNECTIONS CONDA_PY311_CHANNEL PIP_FALLBACK_INDEX_URL
  SMOKE_STEPS FULL_STEPS
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --profile|--baseline)
      PROFILE="$2"; shift 2 ;;
    --bench|--benchmark)
      BENCH="$2"; shift 2 ;;
    --smoke)
      MODE=smoke; TRAIN=1; shift ;;
    --full)
      MODE=full; TRAIN=1; shift ;;
    --prepare-only|--no-train|--skip-train)
      MODE=prepare; TRAIN=0; shift ;;
    --dry-run)
      DRY_RUN=1; shift ;;
    --tmux)
      TMUX_SESSION="$2"; shift 2 ;;
    -h|--help)
      usage; exit 0 ;;
    *)
      echo "unknown arg: $1" >&2
      usage
      exit 2 ;;
  esac
done

if [[ "$PROFILE:$BENCH" != "dreamtacvla:manifeel" ]]; then
  cat >&2 <<EOF
Unsupported verified combo: $PROFILE + $BENCH

This script intentionally only deploys the path that has been reproduced on a
fresh server: dreamtacvla + manifeel. Do not use the old six wrapper scripts as
proof that every matrix cell is deployable.
EOF
  exit 3
fi

if [[ -n "$TMUX_SESSION" && -z "${TMUX:-}" ]]; then
  if command -v tmux >/dev/null 2>&1; then
    mode_arg="--$MODE"
    if [[ "$MODE" == "prepare" ]]; then
      mode_arg="--prepare-only"
    fi
    tmux_args=(--profile "$PROFILE" --bench "$BENCH" "$mode_arg")
    if [[ "$DRY_RUN" -eq 1 ]]; then
      tmux_args+=(--dry-run)
    fi
    exec tmux new -s "$TMUX_SESSION" "$0" "${tmux_args[@]}"
  else
    echo "[warn] --tmux requested but tmux is not installed; continuing in the current shell" >&2
  fi
fi

PROJECT_ROOT=${PROJECT_ROOT:-/root/autodl-tmp/vla_sft}
DATA_ROOT=${DATA_ROOT:-$PROJECT_ROOT/data}
CKPT_ROOT=${CKPT_ROOT:-$PROJECT_ROOT/checkpoints}
CODE_ROOT=${CODE_ROOT:-$PROJECT_ROOT/baselines}
VENV_ROOT=${VENV_ROOT:-$PROJECT_ROOT/envs}
OUTPUT_ROOT=${OUTPUT_ROOT:-$PROJECT_ROOT/outputs}
ARCHIVE_ROOT=${ARCHIVE_ROOT:-$PROJECT_ROOT/.hf_archives}
HF_USERNAME=${HF_USERNAME:-poet70}
SMOKE_STEPS=${SMOKE_STEPS:-10}
FULL_STEPS=${FULL_STEPS:-200000}
ARIA2_CONNECTIONS=${ARIA2_CONNECTIONS:-8}
CONDA_PY311_CHANNEL=${CONDA_PY311_CHANNEL:-https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main}
PIP_FALLBACK_INDEX_URL=${PIP_FALLBACK_INDEX_URL:-https://pypi.org/simple}

CODE_REPO="$HF_USERNAME/vla-sft-code-dreamtacvla"
DATA_REPO="$HF_USERNAME/data-manifeel-dtvla"
CKPT_REPO="$HF_USERNAME/vla-ckpt-dreamtacvla-jepa"
CODE_DIR="$CODE_ROOT/dreamtacvla"
DATA_DIR="$DATA_ROOT/manifeel_hdf5"
CKPT_DIR="$CKPT_ROOT/dreamtacvla_jepa"
DATA_ARCHIVE_DIR="$ARCHIVE_ROOT/manifeel_dtvla"
CKPT_ARCHIVE_DIR="$ARCHIVE_ROOT/dreamtacvla_jepa"
VENV_DIR="$VENV_ROOT/dreamtacvla"
LOG="$OUTPUT_ROOT/deploy_dreamtacvla_manifeel_notes.md"

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

log_note() {
  local msg="$1"
  if [[ "$DRY_RUN" -eq 0 ]]; then
    mkdir -p "$OUTPUT_ROOT"
    printf -- "- %s: %s\n" "$(date)" "$msg" >> "$LOG"
  fi
  echo "[note] $msg"
}

init_notes() {
  run mkdir -p "$PROJECT_ROOT" "$DATA_ROOT" "$CKPT_ROOT" "$CODE_ROOT" "$VENV_ROOT" "$OUTPUT_ROOT" "$ARCHIVE_ROOT"
  if [[ "$DRY_RUN" -eq 0 && ! -f "$LOG" ]]; then
    cat > "$LOG" <<EOF
# DreamTacVLA ManiFeel Deployment Notes

Started: $(date)
Host: $(hostname 2>/dev/null || echo unknown)
Project root: $PROJECT_ROOT

## Progress
EOF
  fi
}

setup_shell_env() {
  if [[ "${USE_NETWORK_TURBO:-auto}" != "0" && -f /etc/network_turbo ]]; then
    # shellcheck disable=SC1091
    source /etc/network_turbo >/dev/null 2>&1 || true
  fi

  export PYTHONUTF8=${PYTHONUTF8:-1}
  export PYTHONIOENCODING=${PYTHONIOENCODING:-utf-8}
  if locale -a 2>/dev/null | grep -Eqi '^(C\.utf8|C\.UTF-8)$'; then
    export LANG=C.UTF-8
    export LC_ALL=C.UTF-8
  elif locale -a 2>/dev/null | grep -Eqi '^en_US\.utf8$'; then
    export LANG=en_US.UTF-8
    export LC_ALL=en_US.UTF-8
  else
    export LANG=C
    export LC_ALL=C
  fi

  export HF_HUB_DISABLE_XET=${HF_HUB_DISABLE_XET:-1}
  export HF_XET_HIGH_PERFORMANCE=${HF_XET_HIGH_PERFORMANCE:-0}
  export HF_HUB_ENABLE_HF_TRANSFER=${HF_HUB_ENABLE_HF_TRANSFER:-0}
}

find_bootstrap_python() {
  if command -v python3 >/dev/null 2>&1; then
    command -v python3
  elif [[ -x /root/miniconda3/bin/python ]]; then
    printf '%s\n' /root/miniconda3/bin/python
  elif command -v python >/dev/null 2>&1; then
    command -v python
  else
    return 1
  fi
}

ensure_hf_cli() {
  if [[ -x "$PROJECT_ROOT/bootstrap_env/bin/hf" ]]; then
    export PATH="$PROJECT_ROOT/bootstrap_env/bin:$PATH"
  fi
  if command -v hf >/dev/null 2>&1; then
    hf --version || true
    return 0
  fi

  local py
  py=$(find_bootstrap_python) || {
    echo "No usable Python found for HF CLI bootstrap." >&2
    exit 4
  }

  log_note "hf CLI not found; creating bootstrap_env with $py."
  run "$py" -m venv "$PROJECT_ROOT/bootstrap_env"
  # shellcheck disable=SC1091
  source "$PROJECT_ROOT/bootstrap_env/bin/activate"
  run python -m pip install -U pip setuptools wheel
  run python -m pip install -U huggingface_hub
  deactivate || true
  export PATH="$PROJECT_ROOT/bootstrap_env/bin:$PATH"
  run hf --version
}

ensure_aria2() {
  if command -v aria2c >/dev/null 2>&1; then
    return 0
  fi
  if command -v apt-get >/dev/null 2>&1 && [[ "$(id -u)" -eq 0 ]]; then
    log_note "aria2c not found; installing aria2 with apt-get for resumable HF archive downloads."
    run apt-get update
    run apt-get install -y aria2
  fi
}

hf_url() {
  local repo="$1"
  local repo_type="$2"
  local path="$3"
  if [[ "$repo_type" == "dataset" ]]; then
    printf 'https://huggingface.co/datasets/%s/resolve/main/%s' "$repo" "$path"
  else
    printf 'https://huggingface.co/%s/resolve/main/%s' "$repo" "$path"
  fi
}

download_url() {
  local url="$1"
  local dest="$2"
  local dir base
  dir=$(dirname "$dest")
  base=$(basename "$dest")
  run mkdir -p "$dir"

  if [[ "$DRY_RUN" -eq 1 ]]; then
    echo "+ download $url -> $dest"
    return 0
  fi

  if command -v aria2c >/dev/null 2>&1; then
    if [[ -n "${HF_TOKEN:-}" ]]; then
      echo "+ aria2c -c -x $ARIA2_CONNECTIONS -s $ARIA2_CONNECTIONS -o $base <url> --header 'Authorization: Bearer <redacted>'"
      if [[ "$DRY_RUN" -eq 0 ]]; then
        aria2c -c -x "$ARIA2_CONNECTIONS" -s "$ARIA2_CONNECTIONS" -k 1M \
          --retry-wait=5 --max-tries=0 --summary-interval=30 --file-allocation=none \
          --header "Authorization: Bearer $HF_TOKEN" -d "$dir" -o "$base" "$url"
      fi
    else
      run aria2c -c -x "$ARIA2_CONNECTIONS" -s "$ARIA2_CONNECTIONS" -k 1M \
        --retry-wait=5 --max-tries=0 --summary-interval=30 --file-allocation=none \
        -d "$dir" -o "$base" "$url"
    fi
  elif command -v wget >/dev/null 2>&1; then
    if [[ -n "${HF_TOKEN:-}" ]]; then
      echo "+ wget -c -O $dest <url> --header 'Authorization: Bearer <redacted>'"
      if [[ "$DRY_RUN" -eq 0 ]]; then
        wget -c --tries=0 --timeout=30 --read-timeout=30 --waitretry=5 \
          --header "Authorization: Bearer $HF_TOKEN" -O "$dest" "$url"
      fi
    else
      run wget -c --tries=0 --timeout=30 --read-timeout=30 --waitretry=5 -O "$dest" "$url"
    fi
  else
    echo "Need aria2c or wget to download archive files." >&2
    exit 4
  fi
}

download_repo() {
  local repo="$1"
  local repo_type="$2"
  local dest="$3"
  run mkdir -p "$dest"
  run hf download "$repo" --repo-type "$repo_type" --local-dir "$dest"
}

artifact_dest_ready() {
  local archive_name="$1"
  local dest_dir="$2"
  case "$archive_name" in
    manifeel-dtvla.tar.gz)
      [[ -d "$dest_dir" ]] && find "$dest_dir" -maxdepth 1 -name '*.hdf5' -print -quit | grep -q .
      ;;
    dreamtacvla-jepa.tar.gz)
      [[ -s "$dest_dir/vitl_manifeel_base.pt" ]]
      ;;
    *)
      return 1
      ;;
  esac
}

extract_archive() {
  local repo="$1"
  local repo_type="$2"
  local archive_name="$3"
  local archive_dir="$4"
  local dest_dir="$5"
  local archive="$archive_dir/$archive_name"
  local checksum="$archive_dir/$archive_name.sha256"

  download_url "$(hf_url "$repo" "$repo_type" "$archive_name.sha256")" "$checksum"
  download_url "$(hf_url "$repo" "$repo_type" "$archive_name")" "$archive"

  if [[ "$DRY_RUN" -eq 0 ]]; then
    (cd "$archive_dir" && sha256sum -c "$archive_name.sha256")
    tar -tzf "$archive" >/tmp/"$archive_name".list
    if artifact_dest_ready "$archive_name" "$dest_dir"; then
      echo "[extract] existing destination already looks ready: $dest_dir"
      return 0
    fi
    mkdir -p "$dest_dir"
    find "$dest_dir" -mindepth 1 -maxdepth 1 -exec rm -rf {} +
    tar -xzf "$archive" -C "$dest_dir" --strip-components=1
  else
    echo "+ sha256sum -c $checksum"
    echo "+ tar -tzf $archive"
    echo "+ extract $archive -> $dest_dir"
  fi
}

download_artifacts() {
  log_note "downloading code, ManiFeel-DTVLA data, and JEPA checkpoint."
  download_repo "$CODE_REPO" model "$CODE_DIR"
  extract_archive "$DATA_REPO" dataset manifeel-dtvla.tar.gz "$DATA_ARCHIVE_DIR" "$DATA_DIR"
  extract_archive "$CKPT_REPO" model dreamtacvla-jepa.tar.gz "$CKPT_ARCHIVE_DIR" "$CKPT_DIR"
  log_note "artifacts downloaded and extracted."
}

patch_dreamtacvla_loader() {
  local target="$CODE_DIR/ModelTrain/module/utils.py"
  if [[ "$DRY_RUN" -eq 1 ]]; then
    echo "+ patch DreamTacVLA loader to include tactile_camera_names"
    return 0
  fi
  python - "$target" <<'PY'
import pathlib
import sys

path = pathlib.Path(sys.argv[1])
text = path.read_text()
new = """            image_dict = dict()
            all_camera_names = list(self.camera_names)
            for cam_name in self.tactile_camera_names:
                if cam_name not in all_camera_names:
                    all_camera_names.append(cam_name)
            for cam_name in all_camera_names:
"""
old = """            image_dict = dict()
            for cam_name in self.camera_names:
"""
if new in text:
    print("loader patch already present")
elif old in text:
    backup = path.with_suffix(path.suffix + ".bak_tactile_loader")
    if not backup.exists():
        backup.write_text(text)
    path.write_text(text.replace(old, new, 1))
    print(f"patched loader: {path}")
else:
    raise SystemExit(f"cannot locate loader camera loop in {path}")
PY
  log_note "patched DreamTacVLA loader to read tactile_camera_names in addition to RGB camera_names."
}

conda_create_python311() {
  local base_dir="$VENV_ROOT/.python311_base"
  if [[ -x "$base_dir/bin/python" ]]; then
    VENV_CREATE_PYTHON="$base_dir/bin/python"
    return 0
  fi
  if ! command -v conda >/dev/null 2>&1; then
    return 1
  fi

  run mkdir -p "$VENV_ROOT"
  if [[ "$DRY_RUN" -eq 1 ]]; then
    echo "+ conda create -y --override-channels -c $CONDA_PY311_CHANNEL -p $base_dir python=3.11"
    VENV_CREATE_PYTHON="$base_dir/bin/python"
    return 0
  fi

  if ! conda create -y --override-channels -c "$CONDA_PY311_CHANNEL" -p "$base_dir" python=3.11; then
    log_note "conda create via $CONDA_PY311_CHANNEL failed; retrying with defaults."
    conda create -y --override-channels -c defaults -p "$base_dir" python=3.11
  fi
  VENV_CREATE_PYTHON="$base_dir/bin/python"
}

select_dreamtacvla_python() {
  if [[ -n "${DREAMTACVLA_PYTHON_BIN:-}" ]]; then
    VENV_CREATE_PYTHON="$DREAMTACVLA_PYTHON_BIN"
  elif command -v python3.11 >/dev/null 2>&1; then
    VENV_CREATE_PYTHON=$(command -v python3.11)
  else
    conda_create_python311 || {
      echo "DreamTacVLA needs Python 3.11. Install python3.11 or conda, or set DREAMTACVLA_PYTHON_BIN." >&2
      exit 4
    }
  fi
}

pip_install() {
  local args=("$@")
  if [[ "$DRY_RUN" -eq 1 ]]; then
    echo "+ python -m pip install ${args[*]}"
    return 0
  fi
  if python -m pip install "${args[@]}"; then
    return 0
  fi
  log_note "pip install failed; retrying with fallback index $PIP_FALLBACK_INDEX_URL."
  python -m pip install -i "$PIP_FALLBACK_INDEX_URL" "${args[@]}"
}

prepare_resnet18_cache() {
  local cache_root dest tmp hash
  cache_root="${TORCH_HOME:-$HOME/.cache/torch}/hub/checkpoints"
  dest="$cache_root/resnet18-f37072fd.pth"
  tmp="$dest.tmp"
  if [[ "$DRY_RUN" -eq 1 ]]; then
    echo "+ cache ResNet18 at $dest"
    return 0
  fi
  if [[ -s "$dest" ]]; then
    return 0
  fi
  mkdir -p "$cache_root"
  curl -L -k --fail --connect-timeout 20 --retry 5 --retry-delay 2 \
    -o "$tmp" https://download.pytorch.org/models/resnet18-f37072fd.pth
  hash=$(sha256sum "$tmp" | awk '{print $1}')
  if [[ "$hash" != f37072fd* ]]; then
    rm -f "$tmp"
    echo "ResNet18 checkpoint hash mismatch: $hash" >&2
    exit 6
  fi
  mv "$tmp" "$dest"
}

install_env() {
  select_dreamtacvla_python
  if [[ ! -d "$VENV_DIR" ]]; then
    run "$VENV_CREATE_PYTHON" -m venv "$VENV_DIR"
  fi
  if [[ "$DRY_RUN" -eq 0 ]]; then
    # shellcheck disable=SC1091
    source "$VENV_DIR/bin/activate"
  else
    echo "+ source $VENV_DIR/bin/activate"
  fi
  run python --version
  pip_install -U pip setuptools wheel
  pip_install torch==2.12.1 torchvision==0.27.1
  if [[ -f "$CODE_DIR/requirements.txt" ]]; then
    pip_install -r "$CODE_DIR/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
  log_note "DreamTacVLA Python environment installed at $VENV_DIR."
}

print_training_command() {
  local steps="$1"
  local out_dir="$2"
  cat <<EOF

Run after GPU is enabled:

export PROJECT_ROOT=$PROJECT_ROOT
export DATA_ROOT=\$PROJECT_ROOT/data
export CKPT_ROOT=\$PROJECT_ROOT/checkpoints
export CODE_ROOT=\$PROJECT_ROOT/baselines
export OUTPUT_ROOT=\$PROJECT_ROOT/outputs
export DOBOT_DATA_DIR=\$DATA_ROOT
export CHECKPOINT_DIR=\$CKPT_ROOT
source "\$PROJECT_ROOT/envs/dreamtacvla/bin/activate"
cd "\$CODE_ROOT/dreamtacvla"
python ModelTrain/model_train.py \\
  --ckpt_dir "$out_dir" \\
  --policy_class ACTJEPAAdapter \\
  --task_name manifeel_all \\
  --batch_size 8 \\
  --seed 42 \\
  --num_steps $steps \\
  --lr 1e-5 \\
  --save_every 100 \\
  --enable_hsa \\
  --freeze_clip \\
  --vit_ckpt_path "\$CKPT_ROOT/dreamtacvla_jepa/vitl_manifeel_base.pt"
EOF
}

run_training() {
  local steps out_dir
  if [[ "$MODE" == "full" ]]; then
    steps="$FULL_STEPS"
    out_dir="$OUTPUT_ROOT/dreamtacvla_manifeel_full"
  else
    steps="$SMOKE_STEPS"
    out_dir="$OUTPUT_ROOT/dreamtacvla_manifeel_smoke"
  fi

  if [[ "$DRY_RUN" -eq 0 ]]; then
    # shellcheck disable=SC1091
    source "$VENV_DIR/bin/activate"
  else
    echo "+ source $VENV_DIR/bin/activate"
  fi
  export DOBOT_DATA_DIR="$DATA_ROOT"
  export CHECKPOINT_DIR="$CKPT_ROOT"
  export OUTPUT_ROOT="$OUTPUT_ROOT"
  run mkdir -p "$out_dir"
  if [[ "$DRY_RUN" -eq 0 ]]; then
    cd "$CODE_DIR"
  else
    echo "+ cd $CODE_DIR"
  fi
  run python ModelTrain/model_train.py \
    --ckpt_dir "$out_dir" \
    --policy_class ACTJEPAAdapter \
    --task_name manifeel_all \
    --batch_size 8 \
    --seed 42 \
    --num_steps "$steps" \
    --lr 1e-5 \
    --save_every 100 \
    --enable_hsa \
    --freeze_clip \
    --vit_ckpt_path "$CKPT_ROOT/dreamtacvla_jepa/vitl_manifeel_base.pt"
}

main() {
  echo "[deploy] verified combo: $PROFILE + $BENCH"
  echo "[deploy] mode: $MODE"
  echo "[deploy] project root: $PROJECT_ROOT"
  setup_shell_env
  init_notes
  ensure_hf_cli
  ensure_aria2
  download_artifacts
  patch_dreamtacvla_loader
  install_env
  if [[ "$TRAIN" -eq 1 ]]; then
    run_training
  else
    log_note "prepare-only completed; training skipped by user request."
    print_training_command "$SMOKE_STEPS" "$OUTPUT_ROOT/dreamtacvla_manifeel_smoke"
  fi
  log_note "deploy_verified.sh completed."
}

main "$@"