FROM python:3.11-slim-bookworm as base

# Use args
ARG USE_CUDA
ARG USE_CUDA_VER

## Basis ##
ENV ENV=prod \
    PORT=9099 \
    # pass build args to the build
    USE_CUDA_DOCKER=${USE_CUDA} \
    USE_CUDA_DOCKER_VER=${USE_CUDA_VER}


# Install GCC and build tools
RUN apt-get update && \
    apt-get install -y gcc build-essential curl git && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*


WORKDIR /app

# Install Python dependencies
COPY ./requirements.txt .
RUN pip3 install uv && \
    if [ "$USE_CUDA" = "true" ]; then \
    pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/$USE_CUDA_DOCKER_VER --no-cache-dir && \
    uv pip install --system -r requirements.txt --no-cache-dir; \
    else \
    pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu --no-cache-dir && \
    uv pip install --system -r requirements.txt --no-cache-dir; \
    fi

# Copy the application code
COPY . .

# Install Rust compiler and ddtrace which are required for DataDog components
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y

# Set up the Rust environment
ENV PATH="/root/.cargo/bin:${PATH}"
RUN /root/.cargo/bin/rustup default stable

# DEBUG - check that Rust installed correctly
RUN cargo --version

# Set the working directory to the Pipelines app dir
WORKDIR /app

# Install Python dependencies
RUN pip3 install git+https://github.com/DataDog/dd-trace-py.git@main

# Expose the port
ENV HOST="0.0.0.0"
ENV PORT="9099"

ENTRYPOINT [ "bash", "start.sh" ]