# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

# ── Stage 1: Base system ─────────────────────────────────────────────
FROM nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04 AS base

ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=UTC
ENV PIP_NO_CACHE_DIR=1

RUN apt-get update && apt-get install -y \
    python3 python3-pip python3-dev \
    ffmpeg git \
    && rm -rf /var/lib/apt/lists/*

RUN ln -sf /usr/bin/python3 /usr/bin/python

WORKDIR /app

# ── Stage 2: Install asset_harvester wheel + Gradio runtime deps ─────
FROM base AS wheel

COPY dist/asset_harvester-1.0.0-py3-none-any.whl /tmp/
RUN pip install --no-cache-dir \
        '/tmp/asset_harvester-1.0.0-py3-none-any.whl[multiview-diffusion,tokengs,camera-estimator]' \
        'gradio>=5.14.0' spaces \
    && rm /tmp/asset_harvester-1.0.0-py3-none-any.whl

# ── Stage 3: gsplat from source (needs torch already installed) ──────
FROM wheel AS gsplat

ARG GSPLAT_COMMIT=b60e917c95afc449c5be33a634f1f457e116ff5e
ENV TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"
RUN pip install --no-cache-dir --no-build-isolation \
    "git+https://github.com/nerfstudio-project/gsplat.git@${GSPLAT_COMMIT}"

# ── Stage 4: Final image ─────────────────────────────────────────────
FROM gsplat AS final

RUN useradd -m -u 1000 user \
    && mkdir -p /app/checkpoints \
    && chown -R 1000:1000 /app

# HF_TOKEN from build secret (optional — can also be passed at runtime via -e)
RUN --mount=type=secret,id=HF_TOKEN,mode=0444 \
    if [ -f /run/secrets/HF_TOKEN ]; then \
        echo "export HF_TOKEN=$(cat /run/secrets/HF_TOKEN)" > /etc/hf_env; \
    else \
        echo "# no build-time HF_TOKEN; provide via -e HF_TOKEN=..." > /etc/hf_env; \
    fi \
    && chmod +x /etc/hf_env

RUN printf '#!/bin/bash\nsource /etc/hf_env\nexec "$@"\n' > /usr/local/bin/entrypoint.sh \
    && chmod +x /usr/local/bin/entrypoint.sh

COPY --chown=1000:1000 app.py /app/
COPY --chown=1000:1000 examples /app/examples

USER user
EXPOSE 7860

ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]
CMD ["python", "app.py"]