#!/usr/bin/env bash
# Robust Puppeteer rigging pipeline
# - stop on errors
# - safe path checks
# - always copy artifacts into /data/results

set -euo pipefail

echo "[INFO] Starting Puppeteer rigging pipeline..."

# ---------- env / paths ----------
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
# repo 레이아웃이 /app/Puppeteer/{skeleton,skinning,...} 라고 가정
ROOT_DIR="$(dirname "$SCRIPT_DIR")"               # /app/Puppeteer
TMP_DIR="/tmp/puppeteer_run"
WORK_DIR="${TMP_DIR}"                              # 통합 작업 디렉토리
IN_EXAMPLES="${ROOT_DIR}/examples"
OUT_ROOT="${WORK_DIR}/results"
OUT_SKEL="${OUT_ROOT}/skel_results"
OUT_SKEL_FOR_SKIN="${OUT_ROOT}/skeletons"
OUT_SKIN="${OUT_ROOT}/skin_results"
OUT_FINAL="${OUT_ROOT}/final_rigging"
RESULT_DIR="${RESULT_DIR:-/data/results}"          # app.py와 동일 환경변수 사용

mkdir -p "$WORK_DIR" "$OUT_ROOT" "$OUT_SKEL_FOR_SKIN" "$OUT_FINAL" "$RESULT_DIR"

# python 경로 보강 (third_party / third_partys 호환)
export PYTHONPATH="/app:/app/Puppeteer:/app/Puppeteer/third_party:${PYTHONPATH:-}"
[ -d /app/third_partys ] || ln -s /app/Puppeteer/third_party /app/third_partys 2>/dev/null || true
[ -f /app/Puppeteer/third_party/__init__.py ] || touch /app/Puppeteer/third_party/__init__.py

# ---------- skeleton ----------
echo "[INFO] Running skeleton generation..."
cd "${ROOT_DIR}/skeleton"

python demo.py \
  --input_dir "${IN_EXAMPLES}" \
  --pretrained_weights skeleton_ckpts/puppeteer_skeleton_w_diverse_pose.pth \
  --output_dir "${OUT_ROOT}" \
  --save_name skel_results \
  --input_pc_num 8192 \
  --save_render \
  --apply_marching_cubes \
  --joint_token \
  --seq_shuffle

echo "[INFO] Skeleton generation completed."

# skel 결과를 skinning 입력 폴더로 복사 (존재 검증)
echo "[INFO] Preparing skeletons for skinning..."
if [ -d "${OUT_SKEL}" ]; then
  mkdir -p "${OUT_SKEL_FOR_SKIN}"
  shopt -s nullglob
  for f in "${OUT_SKEL}"/*_pred.txt; do
    cp -f "$f" "${OUT_SKEL_FOR_SKIN}/$(basename "${f/_pred.txt/.txt}")"
  done
  shopt -u nullglob
else
  echo "[ERR] ${OUT_SKEL} not found (skeleton step failed?)"
  exit 1
fi
echo "[INFO] Copied rig files to ${OUT_SKEL_FOR_SKIN}"

# ---------- skinning ----------
echo "[INFO] Running skinning..."
cd "${ROOT_DIR}/skinning"

# CUDA_VISIBLE_DEVICES는 Space에서 보통 0 하나만 할당됨
CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" \
torchrun --nproc_per_node=1 --master_port=10009 \
  main.py \
  --num_workers 1 \
  --batch_size 1 \
  --generate \
  --save_skin_npy \
  --pretrained_weights skinning_ckpts/puppeteer_skin_w_diverse_pose_depth1.pth \
  --input_skel_folder "${OUT_SKEL_FOR_SKIN}" \
  --mesh_folder "${IN_EXAMPLES}" \
  --post_filter \
  --depth 1 \
  --save_folder "${OUT_SKIN}"

echo "[INFO] Skinning completed."

# ---------- collect artifacts ----------
echo "[INFO] Collecting final artifacts..."
mkdir -p "${OUT_FINAL}"

# 선호 산출물: output/rigged.glb (있으면 우선 복사)
if [ -f "${WORK_DIR}/output/rigged.glb" ]; then
  cp -f "${WORK_DIR}/output/rigged.glb" "${OUT_FINAL}/rigged.glb"
fi

# skin 결과(glb) 있으면 함께 수집
if [ -d "${OUT_SKIN}/generate" ]; then
  shopt -s nullglob
  cp -f "${OUT_SKIN}/generate/"*.glb "${OUT_FINAL}/" 2>/dev/null || true
  shopt -u nullglob
fi

# ---------- export to /data/results ----------
echo "[INFO] Exporting to ${RESULT_DIR} ..."
mkdir -p "${RESULT_DIR}"
shopt -s nullglob
cp -f "${OUT_FINAL}/"*.glb "${RESULT_DIR}/" 2>/dev/null || true
cp -f "${OUT_FINAL}/"*.gltf "${RESULT_DIR}/" 2>/dev/null || true
shopt -u nullglob

# 결과 검증: 최소 하나라도 존재해야 성공 처리
if compgen -G "${RESULT_DIR}/*.glb" > /dev/null || compgen -G "${RESULT_DIR}/*.gltf" > /dev/null ; then
  echo "[OK] Artifacts saved to ${RESULT_DIR}"
else
  echo "[ERR] No .glb/.gltf produced. Check skeleton/skinning logs."
  exit 2
fi

echo "[INFO] Pipeline finished successfully."