ARG BASE_TAG=main
FROM axolotlai/axolotl-uv:$BASE_TAG

USER root

ENV HF_DATASETS_CACHE="/workspace/data/huggingface-cache/datasets"
ENV HF_HUB_CACHE="/workspace/data/huggingface-cache/hub"
ENV HF_HOME="/workspace/data/huggingface-cache/hub"
ENV HF_HUB_ENABLE_HF_TRANSFER="1"

EXPOSE 8888
EXPOSE 22

COPY scripts/cloud-entrypoint.sh /etc/cloud-entrypoint.sh
COPY scripts/motd /etc/motd

RUN uv pip install jupyterlab notebook ipywidgets && \
    jupyter lab clean
RUN apt update && \
    apt install --yes --no-install-recommends openssh-server tmux iproute2 nvtop && \
    rm -rf /var/cache/apt/archives && \
    rm -rf /var/lib/apt/lists/* && \
    mkdir -p /home/ubuntu/.ssh && \
    chmod 700 /home/ubuntu/.ssh && \
    printf "\n[[ -z \"\$TMUX\" ]] && tty -s && { tmux attach-session -t ssh_tmux || tmux new-session -s ssh_tmux; exit; }\n" >> /home/ubuntu/.bashrc && \
    printf "[ ! -z \"\$TERM\" -a -r /etc/motd ] && cat /etc/motd\n" >> /home/ubuntu/.bashrc && \
    printf "\n[[ -z \"\$AXOLOTL_SKIP_SWITCH\" ]] && exec sudo -u ubuntu AXOLOTL_SKIP_SWITCH=1 -i\n" >> /root/.bashrc && \
    chmod +x /etc/cloud-entrypoint.sh && \
    echo 'set-option -g history-limit 5000' >> /home/ubuntu/.tmux.conf && \
    chown -R ubuntu:ubuntu /home/ubuntu /workspace

# USER ubuntu

ENTRYPOINT ["/etc/cloud-entrypoint.sh"]
CMD ["sleep", "infinity"]