ARG BASE_IMAGE=nvcr.io/nvidia/tensorrt:26.04-py3
FROM ${BASE_IMAGE} AS build

RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential curl ca-certificates pkg-config \
    && rm -rf /var/lib/apt/lists/*
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal
ENV PATH=/root/.cargo/bin:${PATH}

WORKDIR /src
COPY rust-batcher/Cargo.toml ./Cargo.toml
COPY rust-batcher/Cargo.lock ./Cargo.lock
COPY rust-batcher/build.rs ./build.rs
COPY rust-batcher/cpp ./cpp
COPY rust-batcher/src ./src
RUN cargo build --release

FROM ${BASE_IMAGE}

ENV RUST_LOG=info \
    DOC_LAYOUT_ENGINE=/models/pp_doclayout_v3.engine \
    DOC_LAYOUT_SAMPLE_IMAGE=/inputs/sample.png \
    DOC_LAYOUT_MAX_BATCH=8 \
    DOC_LAYOUT_MAX_DELAY_US=1000 \
    DOC_LAYOUT_WORKERS=3 \
    DOC_LAYOUT_QUEUE_CAPACITY=4096 \
    DOC_LAYOUT_MAX_UPLOAD_MB=512 \
    DOC_LAYOUT_PORT=18082

COPY --from=build /src/target/release/doclayout-rust-batcher /usr/local/bin/doclayout-rust-batcher
EXPOSE 18082
ENTRYPOINT ["/usr/local/bin/doclayout-rust-batcher"]