ARG INSTRUCTLAB_IMAGE="quay.io/ai-lab/instructlab-amd:latest"
ARG BASEIMAGE="quay.io/centos-bootc/centos-bootc:stream9"
ARG DRIVER_TOOLKIT_IMAGE="quay.io/ai-lab/nvidia-builder:latest"

FROM ${DRIVER_TOOLKIT_IMAGE} AS builder

COPY repos.d/amdgpu.repo /etc/yum.repos.d/amdgpu.repo
COPY repos.d/RPM-GPG-KEY-AMD-ROCM /etc/pki/rpm-gpg/RPM-GPG-KEY-AMD-ROCM

USER root

RUN rpm --import /etc/pki/rpm-gpg/RPM-GPG-KEY-AMD-ROCM \
    && dnf install -y amdgpu-dkms \
    && dnf clean all

FROM ${BASEIMAGE}

ARG VENDOR=''
LABEL vendor=${VENDOR}
LABEL org.opencontainers.image.vendor=${VENDOR}

RUN --mount=type=bind,from=builder,source=/,destination=/tmp/builder,ro \
    export KERNEL_VERSION=$(rpm -q --qf '%{VERSION}-%{RELEASE}.%{ARCH}' kernel-core) \
    && rm -f /lib/modules/${KERNEL_VERSION}/kernel/drivers/gpu/drm/amd/amdgpu/amdgpu.ko.xz \
    && cp -r /tmp/builder/lib/modules/${KERNEL_VERSION}/extra /lib/modules/${KERNEL_VERSION}/extra \
    && cp -r /tmp/builder/lib/firmware/updates/amdgpu /lib/firmware/amdgpu \
    && depmod ${KERNEL_VERSION}

ARG EXTRA_RPM_PACKAGES=''

COPY repos.d/rocm.repo /etc/yum.repos.d/rocm.repo
COPY repos.d/RPM-GPG-KEY-AMD-ROCM /etc/pki/rpm-gpg/RPM-GPG-KEY-AMD-ROCM

RUN rpm --import /etc/pki/rpm-gpg/RPM-GPG-KEY-AMD-ROCM \
    && mv /etc/selinux /etc/selinux.tmp \ 
    && dnf install -y \
        amd-smi \
        cloud-init \
        git \
        git-lfs \
        pciutils \
        rsync \
        skopeo \
        tmux \
        ${EXTRA_RPM_PACKAGES} \
    && dnf clean all \
    && mv /etc/selinux.tmp /etc/selinux \
    && ln -s ../cloud-init.target /usr/lib/systemd/system/default.target.wants

ARG SSHPUBKEY

# The --build-arg "SSHPUBKEY=$(cat ~/.ssh/id_rsa.pub)" option inserts your
# public key into the image, allowing root access via ssh.
RUN if [ -n "${SSHPUBKEY}" ]; then \
    set -eu; mkdir -p /usr/ssh && \
        echo 'AuthorizedKeysFile /usr/ssh/%u.keys .ssh/authorized_keys .ssh/authorized_keys2' >> /etc/ssh/sshd_config.d/30-auth-system.conf && \
            echo ${SSHPUBKEY} > /usr/ssh/root.keys && chmod 0600 /usr/ssh/root.keys; \
fi

# Setup /usr/lib/containers/storage as an additional store for images.
# Remove once the base images have this set by default.
RUN grep -q /usr/lib/containers/storage /etc/containers/storage.conf || \
    sed -i -e '/additionalimage.*/a "/usr/lib/containers/storage",' \
        /etc/containers/storage.conf

ARG INSTRUCTLAB_IMAGE="quay.io/ai-lab/instructlab-amd:latest"
ARG INSTRUCTLAB_IMAGE_PULL_SECRET="instructlab-amd-pull"

COPY duplicated/ilab-wrapper/ilab /usr/bin/ilab
RUN chmod +x /usr/bin/ilab \
    && sed -i "s%__REPLACE_IMAGE_NAME__%${INSTRUCTLAB_IMAGE}%" /usr/bin/ilab

# Added for running as an OCI Container to prevent Overlay on Overlay issues.
VOLUME /var/lib/containers

RUN --mount=type=secret,id=${INSTRUCTLAB_IMAGE_PULL_SECRET}/.dockerconfigjson \
    if [ -f "/run/.input/instructlab-amd/oci-layout" ]; then \
         IID=$(podman --root /usr/lib/containers/storage --storage-opt overlay.force_mask=shared pull oci:/run/.input/instructlab-amd) && \
         podman --root /usr/lib/containers/storage image tag ${IID} ${INSTRUCTLAB_IMAGE}; \
    elif [ -f "/run/secrets/${INSTRUCTLAB_IMAGE_PULL_SECRET}/.dockerconfigjson" ]; then \
         IID=$(sudo podman --root /usr/lib/containers/storage --storage-opt overlay.force_mask=shared pull --authfile /run/secrets/${INSTRUCTLAB_IMAGE_PULL_SECRET}/.dockerconfigjson ${INSTRUCTLAB_IMAGE}); \
    else \
         IID=$(sudo podman --root /usr/lib/containers/storage --storage-opt overlay.force_mask=shared pull ${INSTRUCTLAB_IMAGE}); \
    fi && \
    chmod -R a+X /usr/lib/containers

COPY containers-storage.conf /etc/skel/.config/containers/storage.conf

RUN podman system reset --force 2>/dev/null

LABEL image_version_id="${IMAGE_VERSION_ID}"