# FoundationPose Dockerfile for Hugging Face Spaces with ZeroGPU
FROM nvidia/cuda:12.1.0-devel-ubuntu22.04

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV CUDA_HOME=/usr/local/cuda
ENV PATH=${CUDA_HOME}/bin:${PATH}
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}

# FoundationPose model configuration
# Override these in Space settings to enable real model
ENV FOUNDATIONPOSE_MODEL_REPO=gpue/foundationpose-weights
ENV USE_HF_WEIGHTS=true
ENV USE_REAL_MODEL=false

# Install system dependencies
RUN apt-get update && apt-get install -y \
    git \
    wget \
    python3.9 \
    python3.9-dev \
    python3-pip \
    libgl1-mesa-glx \
    libglib2.0-0 \
    libeigen3-dev \
    && rm -rf /var/lib/apt/lists/*

# Set Python 3.9 as default
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.9 1

# Upgrade pip
RUN python3 -m pip install --upgrade pip setuptools wheel

# Set working directory
WORKDIR /app

# Copy requirements first for better caching
COPY requirements.txt .

# Install Python dependencies
RUN pip install --no-cache-dir -r requirements.txt

# Install FoundationPose dependencies
RUN pip install --no-cache-dir \
    git+https://github.com/NVlabs/nvdiffrast.git

# Install Kaolin (if using model-free mode)
RUN pip install --no-cache-dir kaolin==0.15.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.0.0_cu118.html

# Install PyTorch3D
RUN pip install --no-cache-dir pytorch3d

# Clone FoundationPose repository
RUN git clone https://github.com/NVlabs/FoundationPose.git /app/FoundationPose

# Build FoundationPose C++ extensions
WORKDIR /app/FoundationPose
RUN bash build_all.sh || echo "Build completed with warnings"

# Copy application files
WORKDIR /app
COPY . .

# Download model weights from Hugging Face model repository
# The download_weights.py script will:
# 1. Check if weights already exist locally
# 2. Download from FOUNDATIONPOSE_MODEL_REPO if USE_HF_WEIGHTS=true
# 3. Fall back to placeholder mode if download fails
RUN echo "Downloading model weights from ${FOUNDATIONPOSE_MODEL_REPO}..." && \
    python3 download_weights.py || \
    echo "⚠️  Weights download failed or skipped. Space will run in placeholder mode."

# Expose port for Gradio
EXPOSE 7860

# Run the application
CMD ["python3", "app.py"]