FROM nvidia/cuda:12.5.0-runtime-ubuntu22.04

# Arguments for user ID and group ID, allowing non-root usage
ARG USER_ID
ARG GROUP_ID

# Set environment variables for non-interactive installation and unbuffered Python output
ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1

# Add a non-root user with specified UID and GID
RUN addgroup --gid $GROUP_ID noroot && \
    adduser --disabled-password --gecos "" --uid $USER_ID --gid $GROUP_ID noroot

# Install system dependencies and cleanup apt cache
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        python3 python3-dev python3-pip python3-setuptools \
        tmux libsm6 libxext6 libxrender-dev \
        git wget unzip libglew-dev patchelf && \
    rm -rf /var/lib/apt/lists/*

# Upgrade pip and install Poetry
RUN python3 -m pip install --no-cache-dir --upgrade pip && \
    python3 -m pip install poetry poethepoet poetry-dotenv-plugin && \
    poetry config virtualenvs.create false && \
    rm -rf /root/.cache/pypoetry

# Set no virtualenv also for the non-root user
USER noroot
RUN poetry config virtualenvs.create false
USER root

# Create a script that wraps poetry run --quiet python3 into python command
RUN echo '#!/bin/sh\npoetry --quiet run python3 "$@"' > /usr/local/bin/python && \
    chmod +x /usr/local/bin/python

# Copy and install dependencies from pyproject.toml
COPY pyproject.toml .
RUN poetry install -vvv --no-interaction --no-root

# Set working directory and configure permissions for the workspace
WORKDIR /workspace
RUN chmod -R a+w /workspace && \
    git config --global --add safe.directory /workspace

# Set default command
CMD ["bash"]