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

COPY ./requirements.txt /code/requirements.txt

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

RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt

# Install Visual Studio Code CLI
RUN curl -Lk 'https://code.visualstudio.com/sha/download?build=stable&os=cli-alpine-x64' \
    --output vscode_cli.tar.gz \
    && tar -xvf vscode_cli.tar.gz \
    && chmod +x ./code \
    && mv code /usr/local/bin/

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

# Change ownership of the /data directory
RUN chown -R user:user /data

# Switch to the non-root user
USER user

# Set working directory
WORKDIR /home/user

# Add local python bin directory to PATH
ENV PATH="/home/user/.local/bin:${PATH}"

# 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

# Override the ENTRYPOINT
ENTRYPOINT []

# Switch to the non-root user
USER user

# Set working directory
WORKDIR /home/user

# Set home to the user's home directory
ENV HOME=/home/user \
    PATH=/home/user/.local/bin:$PATH \
    PYTHONPATH=$HOME/app \
    PYTHONUNBUFFERED=1 \
    GRADIO_ALLOW_FLAGGING=never \
    GRADIO_NUM_PORTS=1 \
    GRADIO_SERVER_NAME=0.0.0.0 \
    GRADIO_THEME=huggingface \
    SYSTEM=spaces

# Copy the current directory contents into the container at $HOME/app setting the owner to the user
COPY --chown=user . $HOME/

# Ensure run.sh is executable
RUN chmod +x $HOME/run.sh

# Set the CMD to run your script
CMD ["/home/user/run.sh"]