################################################################################
#                                                                              #
#       Build GenAI Engine on slim-bookworm images for cpu and gpu             #
#                                                                              #
################################################################################

# TORCH_DEVICE must be either "cpu" or "gpu"
ARG TORCH_DEVICE=cpu

# Preinstall Stage: Install Python dependencies
FROM python:3.12.8-slim-bookworm AS preinstall

# Copy requirements files so this layer can be cached / reused when files are otherwise changed in the repo
COPY pyproject.toml /app/
COPY uv.lock /app/

# Install Python dependencies
RUN pip3 install uv==0.11.7
# Set working directory
WORKDIR /app

ENV PYTHONPATH="/app/src"
ENV PYTHONUNBUFFERED=1
RUN uv export --frozen \
    --no-group dev --no-group performance --no-group linters \
    --no-emit-project \
    -o /tmp/requirements.txt \
    && uv pip install --system --no-cache \
    --extra-index-url https://download.pytorch.org/whl/cpu \
    --index-strategy unsafe-best-match \
    -r /tmp/requirements.txt

# install lsof for healthchecks
RUN apt-get update && apt-get install -y lsof && apt-get clean && rm -rf /var/lib/apt/lists/*


# GPU Install: Install PyTorch for GPU
FROM preinstall AS gpu-install
COPY requirements-gpu.txt /tmp/requirements-torch.txt
RUN uv pip install --system --no-cache --reinstall -r /tmp/requirements-torch.txt


# CPU Install: no-op for now
FROM preinstall AS cpu-install

# UI Build Stage: Build the React SPA
FROM node:20-alpine AS ui-build

# Accept Meticulous tokens as build arguments
ARG METICULOUS_RECORDING_TOKEN
ARG METICULOUS_API_TOKEN

# Accept GitLab Unify Frontend token as build argument
ARG GITLAB_UNIFY_FRONTEND_TOKEN

# Accept client-side analytics / anti-abuse keys as build arguments. These are
# baked into the static bundle at `yarn build` time (Vite inlines them), so they
# must be present here rather than as ECS runtime env vars. All are optional;
# when blank the corresponding feature stays disabled in the UI.
ARG AMPLITUDE_API_KEY
ARG VITE_AMPLITUDE_DEPLOYMENT_KEY
ARG RECAPTCHA_ENTERPRISE_SITE_KEY

WORKDIR /app/ui

# Copy UI source code
COPY ui/ ./

# Install UI dependencies (token only in this RUN, not persisted in image)
RUN if [ -z "${GITLAB_UNIFY_FRONTEND_TOKEN}" ]; then \
      echo "ERROR: GITLAB_UNIFY_FRONTEND_TOKEN build-arg is required for yarn install (private @arthur/* packages)." >&2; \
      exit 1; \
    fi \
    && export GITLAB_UNIFY_FRONTEND_TOKEN="${GITLAB_UNIFY_FRONTEND_TOKEN}" \
    && corepack enable \
    && yarn install --immutable

# Build the UI as static files
ENV METICULOUS_RECORDING_TOKEN=${METICULOUS_RECORDING_TOKEN}
ENV METICULOUS_API_TOKEN=${METICULOUS_API_TOKEN}
ENV AMPLITUDE_API_KEY=${AMPLITUDE_API_KEY}
ENV VITE_AMPLITUDE_DEPLOYMENT_KEY=${VITE_AMPLITUDE_DEPLOYMENT_KEY}
ENV RECAPTCHA_ENTERPRISE_SITE_KEY=${RECAPTCHA_ENTERPRISE_SITE_KEY}
RUN yarn build

# Copy the built static files
RUN cp -r dist /app/ui-dist

# Install Stage: Install GenAI Engine on either CPU Install or GPU Install depending on the TORCH_DEVICE variable
FROM ${TORCH_DEVICE}-install AS install
# Copy backend files
COPY src /app/src

# Copy version file to server directory
COPY version /app/src/

# Copy env file to run directory
COPY .env /app/

# Add telemetry setting based on build arg
ARG ENABLE_TELEMETRY=false
RUN echo "TELEMETRY_ENABLED=${ENABLE_TELEMETRY}" >> /app/.env

# Download AWS RDS global certificate bundle
RUN apt-get update && apt-get install -y curl && \
    curl -o /app/postgres-cert.pem https://truststore.pki.rds.amazonaws.com/global/global-bundle.pem

# Copy alembic files to run directory
COPY alembic /app/alembic
COPY alembic.ini /app/

# Copy built UI from ui-build stage
COPY --from=ui-build /app/ui-dist /app/static

# Create models directory for model downloads
RUN mkdir -p /home/nonroot/models

# Pre-download tiktoken encoding files for airgapped environments
RUN TIKTOKEN_CACHE_DIR=/home/nonroot/tiktoken \
    python3 -c "import os, tiktoken; os.makedirs('/home/nonroot/tiktoken', exist_ok=True); [tiktoken.get_encoding(e) for e in ['cl100k_base', 'p50k_base', 'r50k_base', 'o200k_base']]"


#####################################################################################
#                                                                                   #
#    Copy GenAI Engine from install to distroless image to eliminate image bloat    #
#                                                                                   #
#####################################################################################

# Final Stage(s): Create genai-engine image
FROM gcr.io/distroless/python3-debian12:nonroot AS genai_engine_distroless_base

COPY --from=install /bin/sh /bin/sh
COPY --from=install /bin/bash /bin/bash
COPY --from=install /bin/env /bin/env
COPY --from=install /bin/printenv /bin/printenv
COPY --from=install /usr/bin/sh /usr/bin/sh
COPY --from=install /usr/bin/bash /usr/bin/bash
COPY --from=install /usr/bin/env /usr/bin/env
COPY --from=install /usr/bin/printenv /usr/bin/printenv
COPY --from=install /usr/bin/lsof /usr/bin/lsof
COPY --from=install /usr/lib /usr/lib
COPY --from=install /usr/local/lib/ /usr/local/lib/
COPY --from=install /usr/local/bin/ /usr/local/bin/
COPY --from=install /etc/ld.so.cache /etc/ld.so.cache
COPY --from=install --chown=nonroot:nonroot /app/ /home/nonroot/app/
COPY --from=install --chown=nonroot:nonroot /home/nonroot/models /home/nonroot/models
COPY --from=install --chown=nonroot:nonroot /home/nonroot/tiktoken /home/nonroot/tiktoken

ENV PATH="/usr/local/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:$PATH"
ENV PYTHONPATH="/home/nonroot/app/src"
ENV UV_SYSTEM_PYTHON=1
# Default model storage path for nonroot user
ENV MODEL_STORAGE_PATH="/home/nonroot/models"
ENV TIKTOKEN_CACHE_DIR="/home/nonroot/tiktoken"

# Set working directory (this is where the entrypoint will be run)
WORKDIR /home/nonroot/app

# Expose the necessary ports
EXPOSE 3030

ENTRYPOINT ["bash", "src/docker-entrypoint.sh"]
