FROM ubuntu:22.04 AS root
ENV DEBIAN_FRONTEND=noninteractive \
    TZ=Europe/Paris

RUN apt-get update && apt-get install -y \
    fakeroot \
    wget \
    gnupg \
    git \
    git-lfs \
    ffmpeg \
    libsm6 \
    libxext6 \
    cmake \
    rsync \
    libgl1 \
    pip \
    socat \
    curl && \
    mv /usr/bin/apt-get /usr/bin/.apt-get && \
    echo '#!/usr/bin/env sh\nfakeroot /usr/bin/.apt-get $@' > /usr/bin/apt-get && \
    chmod +x /usr/bin/apt-get && \
    git lfs install && \
    rm -rf /var/lib/apt/lists/*

# AWS Neuron GPG Key and repository setup (as root)
RUN wget -qO /tmp/neuron.asc https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB && \
    gpg --no-default-keyring --keyring /tmp/temp-keyring.gpg --import /tmp/neuron.asc && \
    gpg --no-default-keyring --keyring /tmp/temp-keyring.gpg --export --output /etc/apt/trusted.gpg.d/neuron.gpg && \
    echo "deb https://apt.repos.neuron.amazonaws.com jammy main" > /etc/apt/sources.list.d/neuron.list && \
    rm /tmp/neuron.asc /tmp/temp-keyring.gpg*

# Install AWS Neuron packages (as root)
RUN apt-get update && apt-get install -y \
    aws-neuronx-tools \
    aws-neuronx-collectives \
    aws-neuronx-runtime-lib && \
    rm -rf /var/lib/apt/lists/*

# Install Node.js (as root)
RUN curl -fsSL https://deb.nodesource.com/setup_22.x | bash - && \
    apt-get install -y nodejs && \
    rm -rf /var/lib/apt/lists/* && \
    apt-get clean

# Create user
RUN useradd -m -u 1000 user

FROM scratch AS base
COPY --chown=1000:1000 --from=root / /

USER user
ENV HOME=/home/user \
    PATH=/home/user/.local/bin:/opt/aws/neuron/bin:$PATH \
    DEBIAN_FRONTEND=noninteractive \
    TZ=Europe/Paris \
    LD_LIBRARY_PATH=/opt/aws/neuron/lib:$LD_LIBRARY_PATH

WORKDIR /home/user/app

# Install Python packages as user
RUN pip install --no-cache-dir pip -U && \
    pip install torch==2.8.0 torchvision~=0.23 --index-url https://download.pytorch.org/whl/cpu

# Install optimum-neuron
ARG OPTIMUM_NEURON_VERSION=0.4.1
RUN pip install --extra-index-url https://pip.repos.neuron.amazonaws.com --no-cache-dir optimum-neuron[neuronx]==$OPTIMUM_NEURON_VERSION

# Install base Python packages
RUN pip install --no-cache-dir \
    datasets \
    "huggingface-hub>=0.31.4" \
    "protobuf<4" \
    "click<8.1" \
    "pydantic~=1.0"

# Install requirements from file
RUN --mount=target=/tmp/requirements.txt,source=requirements.txt \
    pip install --no-cache-dir -r /tmp/requirements.txt

# Install Streamlit and Gradio
RUN pip install --no-cache-dir \
    gradio[oauth,mcp]==5.31.0 \
    "uvicorn>=0.14.0" \
    spaces

ENTRYPOINT ["python3", "./app.py"]

FROM base AS pipfreeze
RUN pip freeze > /tmp/freeze.txt

FROM base AS run
COPY --link --chown=1000 ./ /home/user/app
COPY --from=pipfreeze --link --chown=1000 /tmp/freeze.txt /tmp/freeze.txt

ENV PYTHONPATH=$HOME/app \
    PYTHONUNBUFFERED=1 \
    HF_XET_HIGH_PERFORMANCE=1 \
    GRADIO_ALLOW_FLAGGING=never \
    GRADIO_FLAGGING_MODE=never \
    GRADIO_CACHE_EXAMPLES=true \
    GRADIO_NUM_PORTS=1 \
    GRADIO_SERVER_NAME=0.0.0.0 \
    GRADIO_SSR_MODE=True \
    GRADIO_THEME=huggingface \
    TQDM_POSITION=-1 \
    TQDM_MININTERVAL=1 \
    SYSTEM=spaces