# Foxhunt Production Build - Multi-stage with cargo-chef
# GLIBC 2.35 (Ubuntu 22.04), CUDA 12.4.1, cuDNN
# Optimized for BuildKit layer caching

# ============================================================================
# Stage 1: Chef Preparation - Install cargo-chef
# ============================================================================
FROM rust:1.82-slim AS chef

# Install cargo-chef for dependency caching
RUN cargo install cargo-chef --locked

WORKDIR /app

# ============================================================================
# Stage 2: Planner - Generate recipe.json for dependency caching
# ============================================================================
FROM chef AS planner

# Copy entire workspace to analyze dependencies
COPY . .

# Generate recipe.json (dependency manifest)
RUN cargo chef prepare --recipe-path recipe.json

# ============================================================================
# Stage 3: Builder Dependencies - Build dependencies in cached layer
# ============================================================================
FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04 AS builder-deps

# Install system dependencies
RUN apt-get update && apt-get install -y \
    curl \
    git \
    build-essential \
    pkg-config \
    libssl-dev \
    ca-certificates \
    && rm -rf /var/lib/apt/lists/*

# Install Rust stable toolchain
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable
ENV PATH="/root/.cargo/bin:${PATH}"

# Install cargo-chef
RUN cargo install cargo-chef --locked

WORKDIR /app

# Copy recipe.json from planner stage
COPY --from=planner /app/recipe.json recipe.json

# Build dependencies (this layer is cached unless dependencies change)
# Use BuildKit cache mount for cargo registry
# NOTE: We build only the ml package with cuda feature, not the entire workspace
# Set CUDA_COMPUTE_CAP to skip nvidia-smi detection (86 = RTX 30-series, 75 = T4/V100)
ENV CUDA_COMPUTE_CAP=86
RUN --mount=type=cache,target=/usr/local/cargo/registry \
    --mount=type=cache,target=/usr/local/cargo/git \
    cargo chef cook --release --recipe-path recipe.json -p ml --features ml/cuda

# ============================================================================
# Stage 4: Builder Application - Build hyperopt binaries
# ============================================================================
FROM builder-deps AS builder

# Copy entire workspace
COPY . .

# Set SQLX offline mode (no database connection during build)
ENV SQLX_OFFLINE=true

# Set CUDA_COMPUTE_CAP to skip nvidia-smi detection (86 = RTX 30-series, 75 = T4/V100)
ENV CUDA_COMPUTE_CAP=86

# Build all 4 hyperopt binaries + train_dqn + train_ppo_parquet with CUDA support
# Use BuildKit cache mount for cargo registry and target directory
RUN --mount=type=cache,target=/usr/local/cargo/registry \
    --mount=type=cache,target=/usr/local/cargo/git \
    --mount=type=cache,target=/app/target \
    cargo build -p ml --release --features ml/cuda \
        --example hyperopt_mamba2_demo \
        --example hyperopt_dqn_demo \
        --example hyperopt_ppo_demo \
        --example hyperopt_tft_demo \
        --example train_dqn \
        --example train_ppo_parquet && \
    # Copy binaries out of cached target directory to persistent location
    mkdir -p /app/binaries && \
    cp /app/target/release/examples/hyperopt_mamba2_demo /app/binaries/ && \
    cp /app/target/release/examples/hyperopt_dqn_demo /app/binaries/ && \
    cp /app/target/release/examples/hyperopt_ppo_demo /app/binaries/ && \
    cp /app/target/release/examples/hyperopt_tft_demo /app/binaries/ && \
    cp /app/target/release/examples/train_dqn /app/binaries/ && \
    cp /app/target/release/examples/train_ppo_parquet /app/binaries/

# Verify binaries exist and strip debug symbols
RUN ls -lh /app/binaries/ && \
    strip /app/binaries/hyperopt_mamba2_demo && \
    strip /app/binaries/hyperopt_dqn_demo && \
    strip /app/binaries/hyperopt_ppo_demo && \
    strip /app/binaries/hyperopt_tft_demo && \
    strip /app/binaries/train_dqn && \
    strip /app/binaries/train_ppo_parquet

# ============================================================================
# Stage 5: Runtime - Minimal production image with CUDA runtime
# ============================================================================
FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04 AS runtime

# Install minimal runtime dependencies
RUN apt-get update && apt-get install -y \
    ca-certificates \
    libssl3 \
    && rm -rf /var/lib/apt/lists/*

# Create non-root user for security
RUN useradd -m -u 1000 foxhunt

# Set working directory
WORKDIR /workspace

# Copy all 4 hyperopt binaries + train_dqn + train_ppo_parquet from builder stage
COPY --from=builder /app/binaries/hyperopt_mamba2_demo /usr/local/bin/
COPY --from=builder /app/binaries/hyperopt_dqn_demo /usr/local/bin/
COPY --from=builder /app/binaries/hyperopt_ppo_demo /usr/local/bin/
COPY --from=builder /app/binaries/hyperopt_tft_demo /usr/local/bin/
COPY --from=builder /app/binaries/train_dqn /usr/local/bin/
COPY --from=builder /app/binaries/train_ppo_parquet /usr/local/bin/

# Set executable permissions
RUN chmod +x /usr/local/bin/hyperopt_mamba2_demo && \
    chmod +x /usr/local/bin/hyperopt_dqn_demo && \
    chmod +x /usr/local/bin/hyperopt_ppo_demo && \
    chmod +x /usr/local/bin/hyperopt_tft_demo && \
    chmod +x /usr/local/bin/train_dqn && \
    chmod +x /usr/local/bin/train_ppo_parquet

# Verify GLIBC version (Ubuntu 22.04 = GLIBC 2.35)
RUN ldd --version

# Copy entrypoint scripts for self-termination
COPY scripts/entrypoint-generic.sh /entrypoint-generic.sh
COPY scripts/entrypoint-self-terminate.sh /entrypoint-self-terminate.sh
RUN chmod +x /entrypoint-generic.sh /entrypoint-self-terminate.sh

# Install runpodctl for pod self-termination
RUN apt-get update && apt-get install -y wget && \
    wget -q https://github.com/runpod/runpodctl/releases/download/v1.14.3/runpodctl-linux-amd64 -O /usr/local/bin/runpodctl && \
    chmod +x /usr/local/bin/runpodctl && \
    rm -rf /var/lib/apt/lists/*

# Change ownership to non-root user
RUN chown -R foxhunt:foxhunt /workspace

# Switch to non-root user
USER foxhunt

# Set CUDA environment variables
ENV CUDA_HOME=/usr/local/cuda
ENV PATH=${CUDA_HOME}/bin:${PATH}
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}

# Set entrypoint to self-terminating wrapper (prevents pod restart after training)
ENTRYPOINT ["/entrypoint-self-terminate.sh"]

# Metadata labels
ARG GIT_COMMIT=unknown
ARG BUILD_DATE=unknown
ARG CUDA_VERSION=12.4.1
ARG CUDNN_VERSION=9

LABEL org.opencontainers.image.title="Foxhunt Hyperopt Binaries" \
      org.opencontainers.image.description="CUDA-enabled hyperparameter optimization for ML models" \
      org.opencontainers.image.version="${GIT_COMMIT}" \
      org.opencontainers.image.created="${BUILD_DATE}" \
      org.opencontainers.image.vendor="Foxhunt HFT Trading System" \
      foxhunt.cuda.version="${CUDA_VERSION}" \
      foxhunt.cudnn.version="${CUDNN_VERSION}" \
      foxhunt.glibc.version="2.35" \
      foxhunt.binaries="hyperopt_mamba2_demo,hyperopt_dqn_demo,hyperopt_ppo_demo,hyperopt_tft_demo"

# Default command: list available binaries
CMD ["sh", "-c", "echo 'Available binaries:' && ls -lh /usr/local/bin/hyperopt_*"]