#!/bin/bash
# imajin-identity CLI wrapper with CUDA 12 library configuration
#
# This script configures LD_LIBRARY_PATH for CUDA 12 pip packages before
# launching the Python CLI. Required because onnxruntime-gpu 1.23.x needs
# CUDA 12 libraries that are installed via pip but not on the system path.

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SERVICE_DIR="$(dirname "$SCRIPT_DIR")"
VENV_DIR="$SERVICE_DIR/.venv"

# Check venv exists
if [[ ! -d "$VENV_DIR" ]]; then
    echo "Error: Virtual environment not found at $VENV_DIR" >&2
    echo "Run: python -m venv $VENV_DIR && pip install -e ." >&2
    exit 1
fi

# Activate venv FIRST to get correct Python version
source "$VENV_DIR/bin/activate"

# Find Python site-packages using venv Python
PYTHON_VERSION=$(python -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
SITE_PACKAGES="$VENV_DIR/lib64/python${PYTHON_VERSION}/site-packages"

if [[ ! -d "$SITE_PACKAGES" ]]; then
    SITE_PACKAGES="$VENV_DIR/lib/python${PYTHON_VERSION}/site-packages"
fi

# Configure CUDA 12 library paths from pip-installed nvidia packages
NVIDIA_BASE="$SITE_PACKAGES/nvidia"
CUDA_PATHS=""

for pkg in cublas cuda_runtime cudnn cufft curand cusolver cusparse nvjitlink; do
    lib_dir="$NVIDIA_BASE/$pkg/lib"
    if [[ -d "$lib_dir" ]]; then
        if [[ -z "$CUDA_PATHS" ]]; then
            CUDA_PATHS="$lib_dir"
        else
            CUDA_PATHS="$CUDA_PATHS:$lib_dir"
        fi
    fi
done

# Export LD_LIBRARY_PATH with CUDA libraries prepended
if [[ -n "$CUDA_PATHS" ]]; then
    export LD_LIBRARY_PATH="$CUDA_PATHS${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
fi

# Run CLI
exec python -m imajin_identity.cli.main "$@"
