FROM nvidia/cuda:11.8.0-cudnn8-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 configuration
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.10 \
    python3.10-dev \
    python3-pip \
    libgl1-mesa-glx \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender-dev \
    libgomp1 \
    && rm -rf /var/lib/apt/lists/*

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

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

# Set working directory
WORKDIR /app

# Install Python dependencies first (for better Docker layer caching)
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# 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 app.py client.py ./

# Create weights directory
RUN mkdir -p weights

# Download weights if USE_HF_WEIGHTS=true (optional at build time)
# Weights can also be downloaded at runtime
RUN python3 -c "import os; \
from pathlib import Path; \
from huggingface_hub import snapshot_download; \
repo = os.environ.get('FOUNDATIONPOSE_MODEL_REPO', 'gpue/foundationpose-weights'); \
token = os.environ.get('HF_TOKEN'); \
use_hf = os.environ.get('USE_HF_WEIGHTS', 'false').lower() == 'true'; \
use_real = os.environ.get('USE_REAL_MODEL', 'false').lower() == 'true'; \
if use_hf and use_real: \
    print(f'Downloading weights from {repo}...'); \
    snapshot_download(repo_id=repo, local_dir='weights', token=token, repo_type='model'); \
    print('✓ Weights downloaded'); \
else: \
    print('Placeholder mode - skipping weights')" || echo "⚠️  Weight download skipped"

# Expose Gradio port
EXPOSE 7860

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