#!/bin/bash
set -e

# --- Configuration ---
BLENDER_VERSION="4.2.0"
BLENDER_MAJOR_MINOR="4.2" # Corresponds to Blender's internal versioned Python directory
BLENDER_PYTHON_VERSION="python3.11" # Should match the -dev package (e.g., python3.11-dev)
BLENDER_TARBALL="blender-${BLENDER_VERSION}-linux-x64.tar.xz"
BLENDER_URL="https://download.blender.org/release/Blender${BLENDER_MAJOR_MINOR}/blender-${BLENDER_VERSION}-linux-x64.tar.xz"

APP_DIR="/home/user/app"
BLENDER_INSTALL_BASE="${APP_DIR}/blender_installation"
INSTALL_DIR="${BLENDER_INSTALL_BASE}/blender-${BLENDER_VERSION}-linux-x64"
LOCAL_BIN_DIR="${APP_DIR}/local_bin"

BLENDER_PY_EXEC="${INSTALL_DIR}/${BLENDER_MAJOR_MINOR}/python/bin/${BLENDER_PYTHON_VERSION}"

UNIRIG_REQS_FILE_IN_SPACE="${APP_DIR}/unirig_requirements.txt"
UNIRIG_REPO_CLONE_DIR="${APP_DIR}/UniRig"

TORCH_VERSION="2.3.1"
TORCHVISION_VERSION="0.18.1"
TARGET_CUDA_VERSION_SHORT="cu121" # PyTorch 2.3.1 wheels are available for cu121
TORCH_INDEX_URL="https://download.pytorch.org/whl/${TARGET_CUDA_VERSION_SHORT}"

# --- Set Environment Variables for Build ---
export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
export PATH="${CUDA_HOME}/bin:${LOCAL_BIN_DIR}:${PATH}"
export MAX_JOBS=${MAX_JOBS:-4}

PYTHON_INCLUDE_DIR="/usr/include/python${BLENDER_PYTHON_VERSION#python}"
export CPATH="${PYTHON_INCLUDE_DIR}:${CPATH}"
export C_INCLUDE_PATH="${PYTHON_INCLUDE_DIR}:${C_INCLUDE_PATH}"
export CPLUS_INCLUDE_PATH="${PYTHON_INCLUDE_DIR}:${CPLUS_INCLUDE_PATH}"

# TORCH_CUDA_ARCH_LIST might still be useful for other compiled packages if any
export TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"

echo "--- Setup Script Start ---"
echo "Target Blender Installation Directory: ${INSTALL_DIR}"
echo "Blender Python Executable: ${BLENDER_PY_EXEC}"
echo "Using CUDA_HOME=${CUDA_HOME}"
echo "Targeting PyTorch for CUDA: ${TARGET_CUDA_VERSION_SHORT}"
echo "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}"
echo "NOTE: flash-attn installation is now SKIPPED."

# --- Download and Extract Blender ---
mkdir -p "${BLENDER_INSTALL_BASE}"
mkdir -p "${LOCAL_BIN_DIR}"

if [ ! -d "${INSTALL_DIR}" ] || [ -z "$(ls -A "${INSTALL_DIR}")" ]; then
    echo "Blender not found or directory empty at ${INSTALL_DIR}. Proceeding with download and extraction."
    echo "Downloading Blender ${BLENDER_VERSION}..."
    if [ ! -f "/tmp/${BLENDER_TARBALL}" ]; then
        wget -nv -O "/tmp/${BLENDER_TARBALL}" ${BLENDER_URL}
    else
        echo "Blender tarball /tmp/${BLENDER_TARBALL} already downloaded."
    fi

    echo "Extracting Blender to ${BLENDER_INSTALL_BASE}..."
    tar -xJf "/tmp/${BLENDER_TARBALL}" -C "${BLENDER_INSTALL_BASE}"
    if [ -d "${INSTALL_DIR}" ]; then
        echo "Blender extracted successfully to ${INSTALL_DIR}"
    else
        echo "ERROR: Blender extraction failed. Expected: ${INSTALL_DIR}"
        ls -la "${BLENDER_INSTALL_BASE}"
        exit 1
    fi
else
    echo "Blender already appears to be extracted to ${INSTALL_DIR}."
fi
echo "Extraction complete."

if [ -f "${INSTALL_DIR}/blender" ]; then
    echo "Creating local symlink for Blender executable in ${LOCAL_BIN_DIR}..."
    ln -sf "${INSTALL_DIR}/blender" "${LOCAL_BIN_DIR}/blender"
    echo "Local symlink created at ${LOCAL_BIN_DIR}/blender."
else
    echo "WARNING: Blender executable not found at ${INSTALL_DIR}/blender."
fi

# --- Install Dependencies into Blender's Python ---
echo "Installing dependencies into Blender's Python (${BLENDER_PY_EXEC})..."
if [ ! -f "${BLENDER_PY_EXEC}" ]; then
    echo "ERROR: Blender Python executable not found at ${BLENDER_PY_EXEC}!"
    exit 1
fi
if [ ! -f "${UNIRIG_REQS_FILE_IN_SPACE}" ]; then
    echo "ERROR: UniRig requirements file not found at ${UNIRIG_REQS_FILE_IN_SPACE}!"
    exit 1
fi

echo "Upgrading pip for Blender Python..."
"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir --upgrade pip setuptools wheel -vvv

echo "Step 1: Installing PyTorch ${TORCH_VERSION} (for CUDA ${TARGET_CUDA_VERSION_SHORT}) and Torchvision ${TORCHVISION_VERSION}..."
"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
    torch==${TORCH_VERSION} \
    torchvision==${TORCHVISION_VERSION} \
    --index-url ${TORCH_INDEX_URL} -vvv
echo "PyTorch and Torchvision installation attempted."

# Step 2 for flash-attn has been removed.

echo "Step 2 (formerly Step 3): Installing remaining dependencies from ${UNIRIG_REQS_FILE_IN_SPACE}..."
# Ensure flash-attn is REMOVED from unirig_requirements.txt.
# This will install torch-scatter, torch-cluster, spconv, bpy, etc.
"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
    -r "${UNIRIG_REQS_FILE_IN_SPACE}" -vvv

echo "Dependency installation for Blender's Python complete."

# --- FIX: Ensure UniRig/src is treated as a package ---
UNIRIG_SRC_DIR="${UNIRIG_REPO_CLONE_DIR}/src"
INIT_PY_PATH="${UNIRIG_SRC_DIR}/__init__.py"
if [ -d "${UNIRIG_SRC_DIR}" ]; then
  if [ ! -f "${INIT_PY_PATH}" ]; then
    echo "Creating missing __init__.py in ${UNIRIG_SRC_DIR}..."
    touch "${INIT_PY_PATH}"
    echo "__init__.py created."
  else
    echo "${INIT_PY_PATH} already exists."
  fi
else
  echo "WARNING: UniRig src directory not found at ${UNIRIG_SRC_DIR}."
fi

# (Optional) VRM Addon installation
VRM_ADDON_REL_PATH="blender/add-on-vrm-v2.20.77_modified.zip"
ABSOLUTE_ADDON_PATH="${UNIRIG_REPO_CLONE_DIR}/${VRM_ADDON_REL_PATH}"
if [ -f "${ABSOLUTE_ADDON_PATH}" ]; then
    echo "Installing optional VRM addon for Blender..."
    (cd "${UNIRIG_REPO_CLONE_DIR}" && \
     "${BLENDER_PY_EXEC}" -c "import bpy, os; bpy.ops.preferences.addon_install(overwrite=True, filepath=os.path.abspath('${VRM_ADDON_REL_PATH}')); bpy.ops.preferences.addon_enable(module='io_scene_vrm'); print('VRM Addon installed and enabled.')")
    echo "VRM addon installation attempted."
else
    echo "VRM addon zip not found at ${ABSOLUTE_ADDON_PATH}, skipping."
fi

# --- Cleanup ---
echo "Cleaning up downloaded Blender tarball..."
rm -f /tmp/${BLENDER_TARBALL}
echo "Cleanup complete."
echo "Blender setup finished successfully. Blender is in ${INSTALL_DIR}"
echo "--- Setup Script End ---"