####################################
# BASE STAGE
####################################
FROM ubuntu:24.04 AS base

ARG PROXY
ARG TORCH_VERSION

# Env variables
## Proxy
ENV HTTP_PROXY=${PROXY}
ENV HTTPS_PROXY=${PROXY}
ENV NO_PROXY=localhost,127.0.0.1
## Torch Related
ENV TORCH_VERSION=${TORCH_VERSION:-2.7.0}

ENV LIBTORCH_BASE_DIR=/opt/torch
ENV LIBTORCH=${LIBTORCH_BASE_DIR}/libtorch
ENV LD_LIBRARY_PATH=${LIBTORCH}/lib
## Rust Related
ENV PATH="/root/.cargo/bin:/root/.local/bin:${PATH}"
ENV CXXFLAGS="-D_GLIBCXX_USE_CXX11_ABI=1"

# Install essential build dependencies
RUN apt-get update && apt-get install -y \
    build-essential \
    pkg-config  \
    unzip  \
    tar \
    git \
    wget \
    curl \
    libssl-dev \
    libpq-dev\
    protobuf-compiler

# Prepare libtorch
WORKDIR /
RUN mkdir -p ${LIBTORCH_BASE_DIR} && \
    wget -q https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-${TORCH_VERSION}%2Bcpu.zip && \
    unzip /libtorch-cxx11-abi-shared-with-deps-${TORCH_VERSION}+cpu.zip -d ${LIBTORCH_BASE_DIR} && \
    rm /libtorch-cxx11-abi-shared-with-deps-${TORCH_VERSION}+cpu.zip

####################################
# RUST BUILD STAGE
####################################
FROM base AS ek-rust-build

# Prepare rust toolchain
WORKDIR /app
COPY rust-toolchain.toml /app/
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain none && \
    rustup show

# Install migration Tools
RUN cargo install diesel_cli --no-default-features --features "postgres"

# Build Rust
WORKDIR /app
## Build dependencies
RUN cargo new --lib ek-base && \
    cargo new --lib ek-benchmark && \
    cargo new --lib ek-computation && \
    cargo new --lib ek-db && \
    cargo new --lib ek-cli
COPY ./ek-base/Cargo.toml ./ek-base/
COPY ./ek-benchmark/Cargo.toml ./ek-benchmark/
COPY ./ek-computation/Cargo.toml ./ek-computation/
COPY ./ek-cli/Cargo.toml ./ek-cli/
COPY ./ek-db/Cargo.toml ./ek-db/
COPY ./Cargo.toml ./Cargo.lock ./
RUN --mount=type=cache,target=/usr/local/cargo/registry cargo build --release --lib

## Build the rest of the project
COPY ek-base/ ./ek-base/
COPY ek-benchmark/ ./ek-benchmark/
COPY ek-computation/ ./ek-computation/
COPY ek-cli/ ./ek-cli/
COPY ek-db/ ./ek-db/
COPY ek-proto/ ./ek-proto/
RUN --mount=type=cache,target=/usr/local/cargo/registry  <<EOF
    touch ./ek-base/src/lib.rs
    touch ./ek-computation/src/lib.rs
    touch ./ek-benchmark/src/lib.rs
    touch ./ek-db/src/lib.rs
    touch ./ek-cli/src/lib.rs
    cargo build --release
EOF

####################################
# RUNTIME STAGE
####################################
FROM base AS ek-runtime

# Set env
WORKDIR /ek
ENV LIBTORCH=/opt/torch/libtorch
ENV LD_LIBRARY_PATH=$LIBTORCH/lib
# Unset proxy
ENV HTTP_PROXY=
ENV HTTPS_PROXY=
ENV NO_PROXY=

# Copy binary from build stage
COPY --from=ek-rust-build /root/.cargo/bin/diesel /usr/local/bin/
COPY --from=ek-rust-build /app/target/release/ek-cli /usr/local/bin/

# Prepare ek default config
# Create default config with a readable format
RUN mkdir -p /etc/expert-kit && cat > /etc/expert-kit/config.yaml <<EOF
inference:
  instance_name: qwen3_moe_30b_local_test
  model_name: ds-tiny
  hidden_dim: 2048
  intermediate_dim: 768
  
db:
  db_dsn: postgres://dev:dev@localhost:5432/dev
  max_conn_size: 32

weight:
  server:
    addr: http://localhost:6543
  cache:
    Fs:
      path: /ek/weight_cache/

worker:
  id: local_test
  listen: 0.0.0.0
  broadcast: 0.0.0.0
  ports:
    main: 51234

controller:
  listen: 0.0.0.0
  broadcast: localhost
  ports:
    intra: 5001
    inter: 5002
EOF

ENTRYPOINT [ "" ]