# syntax=docker/dockerfile:1.6
#
# This Dockerfile is used to build the Llama Stack container image.
# Example:
# docker build \
#   -f containers/Containerfile \
#   --build-arg DISTRO_NAME=starter \
#   --tag llama-stack:starter .

ARG BASE_IMAGE=python:3.12-slim
FROM ${BASE_IMAGE}

ARG UV_HTTP_TIMEOUT=500
ENV UV_HTTP_TIMEOUT=${UV_HTTP_TIMEOUT}
ENV PYTHONDONTWRITEBYTECODE=1
ENV PIP_DISABLE_PIP_VERSION_CHECK=1
WORKDIR /app

RUN set -eux; \
    if command -v dnf >/dev/null 2>&1; then \
        dnf -y update && \
        dnf install -y iputils git net-tools wget \
            vim-minimal python3.12 python3.12-pip python3.12-wheel \
            python3.12-setuptools python3.12-devel gcc gcc-c++ make && \
        ln -sf /usr/bin/pip3.12 /usr/local/bin/pip && \
        ln -sf /usr/bin/python3.12 /usr/local/bin/python && \
        dnf clean all; \
    elif command -v apt-get >/dev/null 2>&1; then \
        apt-get update && \
        apt-get install -y --no-install-recommends \
            iputils-ping net-tools iproute2 dnsutils telnet \
            curl wget git procps psmisc lsof traceroute bubblewrap \
            gcc g++ && \
        rm -rf /var/lib/apt/lists/*; \
    else \
        echo "Unsupported base image: expected dnf or apt-get" >&2; \
        exit 1; \
    fi

RUN pip install --no-cache-dir uv
ENV UV_SYSTEM_PYTHON=1

ARG INSTALL_MODE="pypi"
ARG LLAMA_STACK_DIR="/workspace"
ARG LLAMA_STACK_CLIENT_DIR=""
ARG PYPI_VERSION=""
ARG TEST_PYPI_VERSION=""
ARG KEEP_WORKSPACE=""
ARG DISTRO_NAME="starter"
ARG RUN_CONFIG_PATH=""

ENV INSTALL_MODE=${INSTALL_MODE}
ENV LLAMA_STACK_DIR=${LLAMA_STACK_DIR}
ENV LLAMA_STACK_CLIENT_DIR=${LLAMA_STACK_CLIENT_DIR}
ENV PYPI_VERSION=${PYPI_VERSION}
ENV TEST_PYPI_VERSION=${TEST_PYPI_VERSION}
ENV KEEP_WORKSPACE=${KEEP_WORKSPACE}
ENV DISTRO_NAME=${DISTRO_NAME}
ENV RUN_CONFIG_PATH=${RUN_CONFIG_PATH}

# Copy the repository so editable installs and run configurations are available.
COPY . /workspace

# Install llama-stack
RUN set -eux; \
    if [ "$INSTALL_MODE" = "editable" ]; then \
        if [ ! -d "$LLAMA_STACK_DIR" ]; then \
            echo "INSTALL_MODE=editable requires LLAMA_STACK_DIR to point to a directory inside the build context" >&2; \
            exit 1; \
        fi; \
        uv pip install --no-cache-dir -e "$LLAMA_STACK_DIR"; \
    elif [ "$INSTALL_MODE" = "test-pypi" ]; then \
        uv pip install --no-cache-dir fastapi libcst; \
        if [ -n "$TEST_PYPI_VERSION" ]; then \
            uv pip install --no-cache-dir --extra-index-url https://test.pypi.org/simple/ --index-strategy unsafe-best-match "llama-stack==$TEST_PYPI_VERSION"; \
        else \
            uv pip install --no-cache-dir --extra-index-url https://test.pypi.org/simple/ --index-strategy unsafe-best-match llama-stack; \
        fi; \
    else \
        if [ -n "$PYPI_VERSION" ]; then \
            uv pip install --no-cache-dir "llama-stack==$PYPI_VERSION"; \
        else \
            uv pip install --no-cache-dir llama-stack; \
        fi; \
    fi;

# Install the client package if it is provided
RUN set -eux; \
    if [ -n "$LLAMA_STACK_CLIENT_DIR" ]; then \
        if [ ! -d "$LLAMA_STACK_CLIENT_DIR" ]; then \
            echo "LLAMA_STACK_CLIENT_DIR is set but $LLAMA_STACK_CLIENT_DIR does not exist" >&2; \
            exit 1; \
        fi; \
        uv pip install --no-cache-dir -e "$LLAMA_STACK_CLIENT_DIR"; \
    fi;

# Install the dependencies for the distribution
RUN set -eux; \
    if [ -z "$DISTRO_NAME" ]; then \
        echo "DISTRO_NAME must be provided" >&2; \
        exit 1; \
    fi; \
    deps="$(llama stack list-deps "$DISTRO_NAME")"; \
    if [ -n "$deps" ]; then \
        printf '%s\n' "$deps" | xargs -L1 uv pip install --no-cache-dir; \
    fi

# Cleanup
RUN set -eux; \
    pip uninstall -y uv; \
    should_remove=1; \
    if [ -n "$KEEP_WORKSPACE" ]; then should_remove=0; fi; \
    if [ "$INSTALL_MODE" = "editable" ]; then should_remove=0; fi; \
    case "$RUN_CONFIG_PATH" in \
        /workspace*) should_remove=0 ;; \
    esac; \
    if [ "$should_remove" -eq 1 ] && [ -d /workspace ]; then rm -rf /workspace; fi

RUN cat <<'EOF' >/usr/local/bin/llama-stack-entrypoint.sh
#!/bin/sh
set -e

if [ -n "$RUN_CONFIG_PATH" ] && [ -f "$RUN_CONFIG_PATH" ]; then
  exec llama stack run "$RUN_CONFIG_PATH" "$@"
fi

if [ -n "$DISTRO_NAME" ]; then
  exec llama stack run "$DISTRO_NAME" "$@"
fi

exec llama stack run "$@"
EOF
RUN chmod +x /usr/local/bin/llama-stack-entrypoint.sh

RUN mkdir -p /.llama /.cache && chmod -R g+rw /app /.llama /.cache

ENTRYPOINT ["/usr/local/bin/llama-stack-entrypoint.sh"]