#!/bin/bash

# Exit immediately if a command exits with a non-zero status.
set -e

# --- Configuration ---
BLENDER_VERSION="4.2.0"
BLENDER_MAJOR_MINOR="4.2" # Matches the version directory structure
BLENDER_PYTHON_VERSION="python3.11" # Blender 4.2 uses Python 3.11
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}"
# Assuming unirig_requirements.txt is in the root directory alongside this script and app.py
UNIRIG_REQS_FILE="unirig_requirements.txt"
# Define the specific torch version and index URL (matching unirig_requirements.txt)
TORCH_VERSION="2.3.1"
TORCHVISION_VERSION="0.18.1"
TORCH_CUDA_SUFFIX="cu121" # CUDA version suffix for wheels
TORCH_INDEX_URL="https://download.pytorch.org/whl/${TORCH_CUDA_SUFFIX}" # Make sure this matches your target CUDA

# PyTorch Geometric wheel URLs (Update these if versions change)
PYG_TORCH_VERSION_SUFFIX="2.3.1" # PyTorch version suffix in PyG URLs
PYG_PYTHON_VERSION="cp311"       # Python version suffix
PYG_CUDA_SUFFIX="cu121"          # CUDA version suffix
# Base URL for the directory containing the wheels
PYG_BASE_URL="https://data.pyg.org/whl/torch-${PYG_TORCH_VERSION_SUFFIX}+${PYG_CUDA_SUFFIX}"
# Correct wheel filenames (checked against the index page)
TORCH_SCATTER_WHL_FILENAME="torch_scatter-2.1.2-${PYG_PYTHON_VERSION}-${PYG_PYTHON_VERSION}-linux_x86_64.whl"
TORCH_CLUSTER_WHL_FILENAME="torch_cluster-1.6.3-${PYG_PYTHON_VERSION}-${PYG_PYTHON_VERSION}-linux_x86_64.whl"
# Construct the full URLs
TORCH_SCATTER_WHL_URL="${PYG_BASE_URL}/${TORCH_SCATTER_WHL_FILENAME}"
TORCH_CLUSTER_WHL_URL="${PYG_BASE_URL}/${TORCH_CLUSTER_WHL_FILENAME}"


# --- Download and Extract Blender ---
echo "Downloading Blender ${BLENDER_VERSION}..."
wget -nv -O /tmp/${BLENDER_TARBALL} ${BLENDER_URL}
echo "Download complete."

echo "Extracting Blender to ${INSTALL_DIR}..."
mkdir -p /opt # Ensure /opt exists
tar -xJf /tmp/${BLENDER_TARBALL} -C /opt
echo "Extraction complete."

# --- Create Blender Symlink ---
echo "Creating symlink for Blender executable..."
ln -sf ${INSTALL_DIR}/blender /usr/local/bin/blender # Use -f to force overwrite if exists
echo "Symlink created."

# --- Install UniRig Dependencies into Blender's Python ---
echo "Installing UniRig Python dependencies into Blender's Python (${BLENDER_PYTHON_VERSION})..."

if [ ! -f "${BLENDER_PY_EXEC}" ]; then
    echo "ERROR: Blender Python executable not found at ${BLENDER_PY_EXEC} after extraction!"
    exit 1
fi

if [ ! -f "${UNIRIG_REQS_FILE}" ]; then
    echo "ERROR: UniRig requirements file not found at ${UNIRIG_REQS_FILE}!"
    echo "Ensure this file exists in the root directory."
    exit 1
fi

# Upgrade pip within Blender's Python environment
echo "Upgrading pip for Blender Python..."
"${BLENDER_PY_EXEC}" -m pip install --upgrade pip setuptools wheel

# --- Install PyTorch and Torchvision FIRST ---
# This is crucial because some packages (like torch-scatter) need torch during their own setup
echo "Installing PyTorch ${TORCH_VERSION} and Torchvision ${TORCHVISION_VERSION} first (CUDA: ${TORCH_CUDA_SUFFIX})..."
"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
    torch==${TORCH_VERSION} \
    torchvision==${TORCHVISION_VERSION} \
    --index-url ${TORCH_INDEX_URL}

echo "PyTorch and Torchvision installation complete."

# --- Install torch-scatter and torch-cluster using direct wheel URLs SECOND ---
# This avoids build issues by using pre-compiled wheels directly.
echo "Installing torch-scatter and torch-cluster from specific wheels..."
echo "Scatter wheel URL: ${TORCH_SCATTER_WHL_URL}"
echo "Cluster wheel URL: ${TORCH_CLUSTER_WHL_URL}"
"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
    "${TORCH_SCATTER_WHL_URL}" \
    "${TORCH_CLUSTER_WHL_URL}"

echo "torch-scatter and torch-cluster installation complete."

# --- Install the rest of the packages from unirig_requirements.txt THIRD ---
# Pip will skip torch, torchvision, torch-scatter, torch-cluster if already installed.
echo "Installing remaining packages from ${UNIRIG_REQS_FILE}..."
"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir -r "${UNIRIG_REQS_FILE}"

echo "UniRig dependency installation for Blender Python complete."

# --- Cleanup ---
echo "Cleaning up downloaded tarball..."
rm /tmp/${BLENDER_TARBALL}
echo "Cleanup complete."

echo "Blender setup finished successfully."