# Start from the TGI base image
FROM ghcr.io/huggingface/text-generation-inference:1.3 as base

# Create a non-root user with UID 1000
RUN useradd -m -u 1000 -s /bin/bash jupyteruser

# Switch to the non-root user
USER jupyteruser

# Set working directory
WORKDIR /home/jupyteruser

# Install JupyterLab and plugins
RUN pip install --user jupyterlab jupyterlab-vim==0.15.1 jupyterlab-vimrc

# Copy any necessary files (if needed)
# COPY --chown=jupyteruser:jupyteruser your-files /home/jupyteruser/your-destination

# AWS Sagemaker compatible image
# Assuming this part remains the same from your original Dockerfile
FROM base as sagemaker

COPY sagemaker-entrypoint.sh entrypoint.sh
RUN chmod +x entrypoint.sh

ENTRYPOINT ["./entrypoint.sh"]

# Final image
FROM base

# Switch to the non-root user
USER jupyteruser

# Set working directory
WORKDIR /home/jupyteruser

# Add JupyterLab entrypoint
# Note: Removed '--allow-root' as we're not running as root
ENTRYPOINT ["jupyter", "lab", "--ip=0.0.0.0", "--NotebookApp.token='' --port 7860"]

# Optional: Set CMD to launch TGI or any other command
#CMD ["text-generation-launcher", "--json-output"]