#!/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"

# ** MODIFIED FOR LOCAL BLENDER INSTALLATION **
# Assuming this script is run from /home/user/app (standard for HF Spaces app.py)
APP_DIR="/home/user/app" # Or use $(pwd) if script is guaranteed to be in app root
BLENDER_INSTALL_BASE="${APP_DIR}/blender_installation" # Base directory for Blender versions
INSTALL_DIR="${BLENDER_INSTALL_BASE}/blender-${BLENDER_VERSION}-linux-x64" # Specific version install
LOCAL_BIN_DIR="${APP_DIR}/local_bin" # For local symlinks, if used

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"

# PyTorch specific versions and URLs for pre-installation
TORCH_VERSION="2.3.1"
TORCHVISION_VERSION="0.18.1"
TARGET_CUDA_VERSION_SHORT="cu121" # Ensure this matches ZeroGPU environment
TORCH_INDEX_URL="https://download.pytorch.org/whl/${TARGET_CUDA_VERSION_SHORT}"
# PYG_FIND_LINKS_URL is used when installing from requirements.txt later
# PYG_FIND_LINKS_URL="https://data.pyg.org/whl/torch-${TORCH_VERSION}+${TARGET_CUDA_VERSION_SHORT}.html"

# --- Set Environment Variables for Build ---
export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda} # Keep for system CUDA if needed by compilations
export PATH="${CUDA_HOME}/bin:${LOCAL_BIN_DIR}:${PATH}" # Add local_bin to PATH early
export MAX_JOBS=${MAX_JOBS:-4} # For compilation processes

PYTHON_INCLUDE_DIR="/usr/include/python${BLENDER_PYTHON_VERSION#python}" # e.g., /usr/include/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 "--- Setup Script Start ---"
echo "Target Blender Installation Directory: ${INSTALL_DIR}"
echo "Blender Python Executable: ${BLENDER_PY_EXEC}"
echo "Using CUDA_HOME=${CUDA_HOME}"
echo "Initial PATH: ${PATH}"
echo "CPATH set to: ${CPATH}"
echo "Expected Python include directory: ${PYTHON_INCLUDE_DIR}"

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 ---
# Create the local installation base directory
mkdir -p "${BLENDER_INSTALL_BASE}"
mkdir -p "${LOCAL_BIN_DIR}" # Ensure local_bin exists

# Check if the *specific version's directory* inside INSTALL_DIR exists and is not empty.
# Blender extracts to a folder like "blender-4.2.0-linux-x64"
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}..."
    # Extracts to a subdirectory like blender-4.2.0-linux-x64 inside BLENDER_INSTALL_BASE
    tar -xJf "/tmp/${BLENDER_TARBALL}" -C "${BLENDER_INSTALL_BASE}"
    # The tarball extracts into a folder named "blender-4.2.0-linux-x64", so INSTALL_DIR should be correct.
    if [ -d "${INSTALL_DIR}" ]; then
        echo "Blender extracted successfully to ${INSTALL_DIR}"
    else
        echo "ERROR: Blender extraction failed or extracted to an unexpected path. Expected: ${INSTALL_DIR}"
        # List contents of BLENDER_INSTALL_BASE for debugging
        ls -la "${BLENDER_INSTALL_BASE}"
        exit 1
    fi
else
    echo "Blender already appears to be extracted to ${INSTALL_DIR}."
fi
echo "Extraction complete."

# --- Create Blender Symlink (Optional - to local_bin) ---
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. Make sure ${LOCAL_BIN_DIR} is in your PATH."
else
    echo "WARNING: Blender executable not found at ${INSTALL_DIR}/blender. Cannot create symlink."
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}!"
    echo "Check Blender extraction and paths."
    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} and Torchvision ${TORCHVISION_VERSION} for Blender's Python..."
# Install torch and torchvision first, as other packages (torch-scatter, etc.) depend on it being present for their setup.
# Ensure your unirig_requirements.txt also lists these versions so pip can resolve correctly later if needed,
# but this direct install ensures they are present for build steps of other packages.
"${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."

echo "Step 2: Installing remaining dependencies from ${UNIRIG_REQS_FILE_IN_SPACE}..."
# Now install the rest of the packages.
# unirig_requirements.txt should contain torch-scatter, torch-cluster, spconv, flash_attn (if needed), bpy, etc.
# with their respective find-links URLs or direct package names.
"${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} 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

# (Optional) VRM Addon installation
VRM_ADDON_REL_PATH="blender/add-on-vrm-v2.20.77_modified.zip" # Relative to UniRig repo root
ABSOLUTE_ADDON_PATH="${UNIRIG_REPO_CLONE_DIR}/${VRM_ADDON_REL_PATH}"
if [ -f "${ABSOLUTE_ADDON_PATH}" ]; then
    echo "Installing optional VRM addon for Blender..."
    # Ensure CWD is UniRig repo for relative path in addon script if any
    (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. Blender is in ${INSTALL_DIR}"
echo "--- Setup Script End ---"