# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

ARG SGLANG_IMAGE_TAG="v0.5.4.post1"
ARG BRANCH_TYPE
ARG CARGO_BUILD_JOBS

FROM scratch AS local_src
COPY . /src

FROM lmsysorg/sglang:${SGLANG_IMAGE_TAG}

WORKDIR /sgl-workspace

# Install jq for JSON processing
RUN apt-get update -y \
    && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
        jq \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

# Install dynamo
# Providing --build-arg BRANCH_TYPE=local will editable install the local dynamo repo
# Providing --build-arg BRANCH_TYPE=remote will editable install the remote dynamo repo
# Default is to install the latest published dynamo version
ARG BRANCH_TYPE
ARG CARGO_BUILD_JOBS

COPY --from=local_src /src /tmp/local_src
RUN if [ "$BRANCH_TYPE" = "local" ]; then \
        cp -r /tmp/local_src /sgl-workspace/dynamo; \
    elif [ "$BRANCH_TYPE" = "remote" ]; then \
        git clone https://github.com/ai-dynamo/dynamo.git /sgl-workspace/dynamo; \
    fi

ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS:-16}

# SGLang does not use a venv in their container
RUN if [ "$BRANCH_TYPE" = "local" ]; then \
    cd dynamo/lib/bindings/python && \
    pip install --break-system-packages maturin && \
    maturin build --release && \
    pip install --break-system-packages target/wheels/*.whl && \
    cd /sgl-workspace/dynamo && \
    pip install --break-system-packages -e . && \
    pip install --break-system-packages --requirement /tmp/local_src/container/deps/requirements.txt ; \
  elif [ "$BRANCH_TYPE" = "remote" ]; then \
    cd dynamo/lib/bindings/python && \
    pip install --break-system-packages maturin && \
    maturin build --release && \
    pip install --break-system-packages target/wheels/*.whl && \
    cd /sgl-workspace/dynamo && \
    pip install --break-system-packages -e . && \
    pip install --break-system-packages --requirement /sgl-workspace/dynamo/container/deps/requirements.txt ; \
  else \
    pip install --break-system-packages ai-dynamo ; \
  fi \
&& rm -rf /tmp/local_src

# Install NATS and ETCD
RUN case "$(uname -m)" in \
      x86_64) ARCH=amd64 ;; \
      aarch64) ARCH=arm64 ;; \
      *) echo "Unsupported architecture: $(uname -m)" && exit 1 ;; \
    esac && \
    wget --tries=3 --waitretry=5 \
      https://github.com/nats-io/nats-server/releases/download/v2.10.28/\
nats-server-v2.10.28-${ARCH}.deb && \
    dpkg -i nats-server-v2.10.28-${ARCH}.deb && rm nats-server-v2.10.28-${ARCH}.deb

ENV ETCD_VERSION="v3.5.21"
RUN case "$(uname -m)" in \
      x86_64) ARCH=amd64 ;; \
      aarch64) ARCH=arm64 ;; \
      *) echo "Unsupported architecture: $(uname -m)" && exit 1 ;; \
    esac && \
    wget --tries=3 --waitretry=5 \
      https://github.com/etcd-io/etcd/releases/download/${ETCD_VERSION}/\
etcd-${ETCD_VERSION}-linux-${ARCH}.tar.gz -O /tmp/etcd.tar.gz && \
    mkdir -p /usr/local/bin/etcd && \
    tar -xzf /tmp/etcd.tar.gz \
        -C /usr/local/bin/etcd --strip-components=1 && \
    rm /tmp/etcd.tar.gz

ENV PATH=/usr/local/bin/etcd:$PATH

# Enable forceful shutdown of inflight requests
ENV SGL_FORCE_SHUTDOWN=1

WORKDIR /sgl-workspace/dynamo/components/backends/sglang