#!/bin/bash
set -e

# --- Configuration ---
BLENDER_VERSION="4.2.0"
BLENDER_MAJOR_MINOR="4.2"
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"
INSTALL_DIR="/opt/blender-${BLENDER_VERSION}-linux-x64"
BLENDER_PY_EXEC="${INSTALL_DIR}/${BLENDER_MAJOR_MINOR}/python/bin/${BLENDER_PYTHON_VERSION}"

# Path to your Gradio Space's unirig_requirements.txt (copy of UniRig's official)
UNIRIG_REQS_FILE_IN_SPACE="/home/user/app/unirig_requirements.txt"
# Path to the cloned UniRig repository in the Space
UNIRIG_REPO_CLONE_DIR="/home/user/app/UniRig" # Standard path

TORCH_VERSION="2.3.1"
TORCHVISION_VERSION="0.18.1"
NUMPY_VERSION="1.26.4"
TARGET_CUDA_VERSION_SHORT="cu121" # Verify ZeroGPU environment.

TORCH_INDEX_URL="https://download.pytorch.org/whl/${TARGET_CUDA_VERSION_SHORT}"
PYG_FIND_LINKS_URL="https://data.pyg.org/whl/torch-${TORCH_VERSION}+${TARGET_CUDA_VERSION_SHORT}.html"
SPCONV_PACKAGE="spconv-${TARGET_CUDA_VERSION_SHORT}"

# --- Set Environment Variables for Build ---
export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
export PATH="${CUDA_HOME}/bin:${PATH}"
export MAX_JOBS=4 # For compilation processes like flash-attn
# ** FIX: Explicitly add Python 3.11 include path for C/C++ header files **
# This assumes python3.11-dev is installed via packages.txt, providing headers here.
PYTHON_INCLUDE_DIR="/usr/include/python${BLENDER_PYTHON_VERSION#python}" # Extracts "3.11" from "python3.11"
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}"

echo "Using CUDA_HOME=${CUDA_HOME}"
echo "Updated PATH: ${PATH}"
echo "CPATH set to: ${CPATH}"
echo "C_INCLUDE_PATH set to: ${C_INCLUDE_PATH}"
echo "CPLUS_INCLUDE_PATH set to: ${CPLUS_INCLUDE_PATH}"
echo "Expected Python include directory: ${PYTHON_INCLUDE_DIR}"
# Check if the Python.h file actually exists where we expect it
if [ -f "${PYTHON_INCLUDE_DIR}/Python.h" ]; then
    echo "Found Python.h at ${PYTHON_INCLUDE_DIR}/Python.h"
else
    echo "WARNING: Python.h NOT FOUND at ${PYTHON_INCLUDE_DIR}/Python.h. Ensure python${BLENDER_PYTHON_VERSION#python}-dev is in packages.txt and installed."
fi


# --- Download and Extract Blender ---
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 ${INSTALL_DIR}..."
if [ ! -d "${INSTALL_DIR}/${BLENDER_MAJOR_MINOR}" ]; then
    mkdir -p /opt
    tar -xJf /tmp/${BLENDER_TARBALL} -C /opt
else
    echo "Blender already extracted to ${INSTALL_DIR}."
fi
echo "Extraction complete."

# --- Create Blender Symlink ---
echo "Creating symlink for Blender executable..."
ln -sf ${INSTALL_DIR}/blender /usr/local/bin/blender
echo "Symlink created."

# --- 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..."
# ** ADDED -vvv for verbosity **
"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir --upgrade pip setuptools wheel -vvv

echo "1. Installing PyTorch ${TORCH_VERSION} and Torchvision ${TORCHVISION_VERSION}..."
# ** ADDED -vvv for verbosity **
"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
    torch==${TORCH_VERSION} \
    torchvision==${TORCHVISION_VERSION} \
    --index-url ${TORCH_INDEX_URL} -vvv

echo "2. Installing dependencies from ${UNIRIG_REQS_FILE_IN_SPACE} (includes bpy, flash_attn)..."
# flash-attn compilation might also benefit from CPATH
# ** ADDED -vvv for verbosity **
"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir -r "${UNIRIG_REQS_FILE_IN_SPACE}" -vvv

echo "3. Installing ${SPCONV_PACKAGE}..."
# ** ADDED -vvv for verbosity **
"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir ${SPCONV_PACKAGE} -vvv

echo "4. Installing torch-scatter and torch-cluster..."
# The CPATH, C_INCLUDE_PATH, CPLUS_INCLUDE_PATH set above should help this step
# ** ADDED -vvv for verbosity **
"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
    torch-scatter \
    torch-cluster \
    -f "${PYG_FIND_LINKS_URL}" -vvv

echo "5. Installing numpy==${NUMPY_VERSION}..."
# ** ADDED -vvv for verbosity **
"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir numpy==${NUMPY_VERSION} -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} to make it a package..."
    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}. Cannot create __init__.py."
fi
# --- END FIX ---

# (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 addon installation."
fi

# --- Cleanup ---
echo "Cleaning up downloaded Blender tarball..."
rm -f /tmp/${BLENDER_TARBALL}
echo "Cleanup complete."
echo "Blender setup finished successfully."