# Copyright 2025 DataRobot, Inc. and its affiliates.
# All rights reserved.
# DataRobot, Inc. Confidential.
# This is unpublished proprietary source code of DataRobot, Inc.
# and its affiliates.
# The copyright notice above does not evidence any actual or intended
# publication of such source code.

ARG WORKDIR=/etc/system/kernel
ARG AGENTDIR=/etc/system/kernel/agent
ARG UV_CACHE_DIR=/data/.uv_cache
ARG PULUMI_PLUGIN_CACHE_DIR=/data/.pulumi_plugin_cache
ARG VENV_PATH=${WORKDIR}/.venv

ARG UNAME=notebooks
ARG UID=10101
ARG GID=10101

# This reproduces a open source variant of https://images.chainguard.dev/directory/image/python
# hadolint ignore=DL3007
FROM cgr.dev/chainguard/wolfi-base:latest AS chainguard_python_dev

ARG PYTHON_VERSION=3.11
SHELL ["/bin/sh", "-o", "pipefail", "-c"]
RUN echo "[Python $PYTHON_VERSION]"

ENV PATH="${PATH}:/sbin"
RUN apk update && apk upgrade
# Wolfi python image
# hadolint ignore=DL3018,DL3059
RUN apk add --no-cache ca-certificates-bundle gdbm glibc glibc-locale-posix \
ld-linux libbz2-1 libcrypto3 libexpat1 libffi libgcc libssl3 libstdc++ \
mpdecimal ncurses ncurses-terminfo readline sqlite-libs \
wolfi-baselayout xz zlib
# Wolfi python dev image
# hadolint ignore=DL3018,DL3059
RUN apk add --no-cache apk-tools  bash  binutils  build-base  busybox  cyrus-sasl  \
gcc  git  glibc-dev  gmp  heimdal-libs  isl  keyutils-libs  \
krb5-conf  krb5-libs  libatomic  libbrotlicommon1  libbrotlidec1  \
libcom_err  libcrypt1  libcurl-openssl4  libgo  libgomp  libidn2  \
libldap  libnghttp2-14  libpcre2-8-0  libpsl  libquadmath  libstdc++-dev  \
libunistring  libverto  libxcrypt  libxcrypt-dev  libzstd1  \
linux-headers  make  mpc  mpfr  nss-db  nss-hesiod  \
openssf-compiler-options  pkgconf  posix-cc-wrappers  wget
# Python
# hadolint ignore=DL3018,DL3059
RUN apk add --no-cache \
    python-$PYTHON_VERSION \
    python-$PYTHON_VERSION-base \
    py$PYTHON_VERSION-pip \
    py3-pip-wheel
# Python dev
# hadolint ignore=DL3018,DL3059
RUN apk add --no-cache python-$PYTHON_VERSION-dev \
    python-$PYTHON_VERSION-base-dev \
    py$PYTHON_VERSION-pip-base \
    py$PYTHON_VERSION-setuptools

FROM chainguard_python_dev AS base

ARG UNAME
ARG UID
ARG GID
ARG WORKDIR
ARG AGENTDIR
ARG PULUMI_PLUGIN_CACHE_DIR
ARG UV_CACHE_DIR
ARG VENV_PATH

USER root

SHELL ["/bin/bash", "-o", "pipefail", "-c"]
# hadolint ignore=DL3018,DL3059
RUN apk add --no-cache uv graphviz openblas openssh-server gzip zip unzip curl \
  openjdk-11 vim nano procps tzdata
# hadolint ignore=DL3018,DL3059
RUN apk add --no-cache rust sqlite task pulumi pulumi-language-python

ENV PYTHONUNBUFFERED=1 \
    PYTHONDONTWRITEBYTECODE=1 \
    VENV_PATH=${VENV_PATH} \
    PIP_NO_CACHE_DIR=1 \
    UV_CACHE_DIR=${UV_CACHE_DIR} \
    UV_LINK_MODE=copy \
    PULUMI_PLUGIN_CACHE_DIR=${PULUMI_PLUGIN_CACHE_DIR} \
    NOTEBOOKS_KERNEL="python" \
    DEEPEVAL_HOME=/tmp/.deepeval \
    DEEPEVAL_TELEMETRY_OPT_OUT="YES"

ENV PATH="$VENV_PATH/bin:$PATH" \
    PYTHONPATH="/home/notebooks/.ipython/extensions:/home/notebooks/storage"

# hadolint ignore=SC1091
RUN uv venv ${VENV_PATH} && mkdir -p ${UV_CACHE_DIR} ${PULUMI_PLUGIN_CACHE_DIR} /tmp/.cache_venv
WORKDIR ${WORKDIR}

COPY ./agent/agent.py ./agent/cgroup_watchers.py ${AGENTDIR}/
COPY ./jupyter_kernel_gateway_config.py ./start_server.sh ${WORKDIR}/
COPY ./ipython_config.py /etc/ipython/
COPY ./extensions /etc/ipython/extensions

# Adding SSHD requirements
COPY ./sshd_config /etc/ssh/
RUN cp -a /etc/ssh /etc/ssh.cache && rm -rf /var/lib/apt/lists/*
RUN mkdir /etc/authorized_keys

# Custom user to run the image from
RUN addgroup -g "$GID" "$UNAME" && \
    adduser -D -u "$UID" -G "$UNAME" -s /bin/bash "$UNAME"

# Prompt customizations
COPY ./setup-prompt.sh /etc/profile.d/setup-prompt.sh

# additional setup scripts
COPY ./setup-ssh.sh ./common-user-limits.sh ./setup-venv.sh ${WORKDIR}/

# Adding SSHD requirements
RUN chown -R $UNAME:$UNAME ${WORKDIR} ${VENV_PATH} ${UV_CACHE_DIR} ${PULUMI_PLUGIN_CACHE_DIR} /tmp/.cache_venv /home/notebooks /etc/ssh /etc/authorized_keys \
  # sshd prep
  && touch /etc/profile.d/notebooks-load-env.sh \
  && chown -R $UNAME:$UNAME /etc/profile.d/notebooks-load-env.sh \
  # Limit max processes
  && touch /etc/profile.d/bash-profile-load.sh \
  && chown -R $UNAME:$UNAME /etc/profile.d/bash-profile-load.sh

USER $UNAME

# Jupyter Gateway port
EXPOSE 8888
# sshd port
EXPOSE 22

FROM base AS builder
# this stage has only bare minimal of dependencies installed to optimize build time for the local development

ENV ANNOY_COMPILER_ARGS="-D_CRT_SECURE_NO_WARNINGS,-DANNOYLIB_MULTITHREADED_BUILD,-march=x86-64"

ARG WORKDIR
ARG UV_CACHE_DIR
ARG PULUMI_PLUGIN_CACHE_DIR
ARG VENV_PATH
ARG TARGETOS
ARG TARGETARCH

COPY --chown=${UNAME}:${UNAME} ./requirements.txt ${WORKDIR}/
# Workspace files needed for `uv sync` to populate cache with lock-file-compatible URL keys
COPY --chown=${UNAME}:${UNAME} ./app_components/ ${WORKDIR}/app_components/

ENV UV_CACHE_DIR=${UV_CACHE_DIR} \
    UV_LINK_MODE=copy \
    PULUMI_PLUGIN_CACHE_DIR=${PULUMI_PLUGIN_CACHE_DIR}

# Use `uv sync` (not `uv pip install`) to populate the UV cache into a temp venv.
# This ensures cache entries are keyed by the exact URLs in uv.lock files,
# so that in codespace or deployment `uv sync --frozen` can find every package via URL lookup.
# Cache is populated from runtime and dev dependencies of all agents, infra projects, etc.
# UV_PROJECT_ENVIRONMENT redirects the venv away from VENV_PATH (cache-only goal).
# --no-install-workspace: downloads and caches all deps without needing package source code.
# Permission on cache dir are modified to allow using this cache by another user (deployment and
# codespace run under different users).
# hadolint ignore=DL3003
RUN for component_dir in "${WORKDIR}"/app_components/*/; do \
        UV_PROJECT_ENVIRONMENT="${component_dir}.venv" \
        uv sync --project "${component_dir}" --frozen --no-install-workspace; \
    done && \
    chmod -R a+rwX ${UV_CACHE_DIR} && \
    rm -rf ${WORKDIR}/app_components

# Install runtime dependencies into the baked venv. Packages are served from the
# cache populated above, so no network access is needed here.
# hadolint ignore=DL3013, SC1091
RUN source ${VENV_PATH}/bin/activate && \
    uv pip install -r ${WORKDIR}/requirements.txt && rm ${WORKDIR}/requirements.txt

# Download pulumi plugin into cache directory.
RUN wget -q -P ${PULUMI_PLUGIN_CACHE_DIR} https://github.com/datarobot-community/pulumi-datarobot/releases/download/v0.10.30/pulumi-resource-datarobot-v0.10.30-$TARGETOS-$TARGETARCH.tar.gz && \
    wget -q -P ${PULUMI_PLUGIN_CACHE_DIR} https://github.com/pulumi/pulumi-command/releases/download/v1.2.1/pulumi-resource-command-v1.2.1-$TARGETOS-$TARGETARCH.tar.gz

# Copy agent runtime into work directory
COPY ./run_agent.py ${WORKDIR}/

RUN rm ${VENV_PATH}/share/jupyter/kernels/python3/kernel.json && chmod a+x ${WORKDIR}/start_server.sh
COPY ./kernel.json ${VENV_PATH}/share/jupyter/kernels/python3/

# Monitoring agent port
EXPOSE 8889

FROM base AS kernel
# this stage is what actually going to be run as kernel image and it's clean from all build junks

ARG UNAME
ARG WORKDIR
ARG PULUMI_PLUGIN_CACHE_DIR
ARG UV_CACHE_DIR
ARG GIT_COMMIT

LABEL com.datarobot.repo-name="notebooks"
LABEL com.datarobot.repo-sha=$GIT_COMMIT

USER root
RUN chown -R $UNAME:$UNAME ${WORKDIR} /home/notebooks && \
    # Ensure directory for dynamic venv creation is writable in LRS.
    mkdir -p /opt/venv && chmod a+rwx /opt/venv
USER $UNAME

COPY --from=builder --chown=$UNAME $WORKDIR $WORKDIR
COPY --from=builder --chown=$UNAME $UV_CACHE_DIR $UV_CACHE_DIR
COPY --from=builder --chown=$UNAME $PULUMI_PLUGIN_CACHE_DIR $PULUMI_PLUGIN_CACHE_DIR
RUN chmod a+rwx ${UV_CACHE_DIR}

# This is required for custom models to work with this image
COPY ./start_server_drum.sh /opt/code/start_server.sh
ENV VENV_DIR=/opt/venv \
    UV_CACHE_DIR=${UV_CACHE_DIR} \
    UV_LINK_MODE=copy \
    UV_FROZEN=1 \
    PULUMI_SKIP_UPDATE_CHECK=1 \
    PULUMI_PLUGIN_CACHE_DIR=${PULUMI_PLUGIN_CACHE_DIR} \
    PULUMI_DISABLE_AUTOMATIC_PLUGIN_ACQUISITION=1

ENV HOME=/opt CODE_DIR=/opt/code ADDRESS=0.0.0.0:8080
