# Base image with CUDA compiler tools (needed for C++ extensions)
FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04

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}

# Only build for T4 (7.5) - reduces compilation memory by 50%
ENV TORCH_CUDA_ARCH_LIST="7.5"

# Install minimal runtime dependencies
# Remove problematic CUDA repo and install packages
RUN rm -f /etc/apt/sources.list.d/cuda*.list /etc/apt/sources.list.d/*.list && \
    apt-get update && apt-get install -y --no-install-recommends --allow-unauthenticated \
    ca-certificates \
    && apt-get clean && rm -rf /var/lib/apt/lists/* && \
    apt-get update && apt-get install -y --no-install-recommends \
    python3.10 \
    python3-pip \
    git \
    libgl1 \
    libglib2.0-0 \
    libgomp1 \
    && rm -rf /var/lib/apt/lists/* \
    && apt-get clean

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

# Upgrade pip
RUN python3 -m pip install --no-cache-dir --upgrade pip

WORKDIR /app

# Install PyTorch (smallest CUDA 11.8 build)
RUN pip install --no-cache-dir torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu118

# Install only essential requirements
# Pin NumPy to 1.x for CUDA extension compatibility
RUN pip install --no-cache-dir \
    "numpy<2" \
    gradio>=4.0.0 \
    opencv-python-headless>=4.8.0 \
    Pillow>=10.0.0 \
    huggingface-hub>=0.20.0 \
    && pip cache purge

# Install build dependencies (keep them for faster HuggingFace builds)
# Install BEFORE nvdiffrast because it needs python3.10-dev
RUN apt-get update && apt-get install -y --no-install-recommends \
    cmake \
    build-essential \
    ninja-build \
    libeigen3-dev \
    python3.10-dev \
    libboost-system-dev \
    libboost-program-options-dev \
    pybind11-dev \
    && rm -rf /var/lib/apt/lists/*

# Install FoundationPose dependencies
RUN pip install --no-cache-dir \
    trimesh==4.2.2 \
    scipy==1.12.0 \
    scikit-image==0.22.0 \
    kornia==0.7.2 \
    einops==0.7.0 \
    timm==0.9.16 \
    transformations==2024.6.1 \
    pyyaml==6.0.1 \
    joblib==1.4.0 \
    psutil==6.1.1 \
    && pip cache purge

# Note: nvdiffrast will be built in final Dockerfile on HuggingFace (needs GPU)

# Clone FoundationPose
RUN git clone --depth 1 https://github.com/NVlabs/FoundationPose.git /app/FoundationPose && \
    cd /app/FoundationPose/bundlesdf/mycuda && \
    sed -i 's/-std=c++14/-std=c++17/g' setup.py

# Build mycpp (non-GPU C++ code - can be built without GPU)
WORKDIR /app/FoundationPose
RUN cd mycpp && mkdir -p build && cd build && cmake .. && make

# Download model weights (246MB)
WORKDIR /app
RUN python3 -c "from huggingface_hub import snapshot_download; \
snapshot_download(repo_id='gpue/foundationpose-weights', local_dir='weights', repo_type='model')"

# Note: Application files (app.py, client.py, estimator.py) are copied in main Dockerfile
# This allows updates without rebuilding the entire base image

EXPOSE 7860