# =============================================================================
# RUNPOD DEPLOYMENT DOCKERFILE - MULTI-STAGE BUILD WITH CUDA 12.1
# =============================================================================
# Purpose: Compile training binaries inside Docker with CUDA 12.1 for perfect compatibility
# Architecture: 2-stage build (builder + runtime)
# Build time: ~25 minutes first build, ~5 minutes incremental
# Final image size: ~8.4GB (includes CUDA 12.1 runtime)
#
# WHY THIS APPROACH:
# - Local CUDA 12.9 binaries don't work with CUDA 12.1 runtime (libcublas.so.13 mismatch)
# - Building inside Docker ensures perfect CUDA version alignment
# - Multi-stage build keeps final image small (no build tools in runtime)
# =============================================================================

# =============================================================================
# STAGE 1: BUILDER - Compile binaries with CUDA 12.1
# =============================================================================
FROM nvidia/cuda:12.1.0-devel-ubuntu22.04 AS builder

ENV DEBIAN_FRONTEND=noninteractive

# Install build dependencies
RUN apt-get update && apt-get install -y \
    curl \
    build-essential \
    pkg-config \
    libssl-dev \
    ca-certificates \
    git \
    && 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}"

# Set CUDA environment for cargo build
ENV CUDA_HOME=/usr/local/cuda-12.1
ENV PATH="${CUDA_HOME}/bin:${PATH}"
ENV LD_LIBRARY_PATH="${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}"
ENV CUDA_COMPUTE_CAP="75,80,86,89,90"

# Copy source code
WORKDIR /build
COPY . .

# Build training binaries with CUDA 12.1
# Note: Using --release for production performance
# Note: Building only ml examples (training binaries)
RUN cargo build --release --features cuda -p ml --example train_tft_parquet && \
    cargo build --release --features cuda -p ml --example train_mamba2_dbn && \
    cargo build --release --features cuda -p ml --example train_dqn && \
    cargo build --release --features cuda -p ml --example train_ppo

# Verify binaries exist and check CUDA dependencies
RUN ls -lh /build/target/release/examples/ && \
    ldd /build/target/release/examples/train_tft_parquet | grep cuda || true

# =============================================================================
# STAGE 2: RUNTIME - Minimal image with CUDA 12.1 runtime
# =============================================================================
FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04 AS runtime

ENV DEBIAN_FRONTEND=noninteractive

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

# 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}"
ENV NVIDIA_VISIBLE_DEVICES=all
ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
ENV CUDA_VISIBLE_DEVICES=0

# Rust runtime environment
ENV RUST_BACKTRACE=1
ENV RUST_LOG=info

# Copy compiled binaries from builder stage
COPY --from=builder /build/target/release/examples/train_tft_parquet /usr/local/bin/
COPY --from=builder /build/target/release/examples/train_mamba2_dbn /usr/local/bin/
COPY --from=builder /build/target/release/examples/train_dqn /usr/local/bin/
COPY --from=builder /build/target/release/examples/train_ppo /usr/local/bin/

# Copy entrypoint script
COPY entrypoint.sh /entrypoint.sh
RUN chmod +x /entrypoint.sh

# Create workspace directory
WORKDIR /workspace

# Health check
HEALTHCHECK --interval=60s --timeout=10s --start-period=30s --retries=3 \
    CMD nvidia-smi || exit 1

# Entrypoint
ENTRYPOINT ["/entrypoint.sh"]
CMD ["--help"]

# =============================================================================
# BUILD INSTRUCTIONS
# =============================================================================
#
# 1. BUILD IMAGE (first build ~25 min, incremental ~5 min):
#    docker build -f Dockerfile.runpod.builder -t jgrusewski/foxhunt:cuda12-built .
#
# 2. TEST LOCALLY (if you have NVIDIA Docker runtime):
#    docker run --gpus all -it jgrusewski/foxhunt:cuda12-built \
#      /usr/local/bin/train_tft_parquet --help
#
# 3. PUSH TO DOCKER HUB:
#    docker login
#    docker push jgrusewski/foxhunt:cuda12-built
#
# 4. DEPLOY ON RUNPOD:
#    - GPU: RTX A4000 or better (16GB VRAM)
#    - Docker Image: jgrusewski/foxhunt:cuda12-built
#    - Volume Mount: /runpod-volume (for test data and model outputs)
#    - Docker Start Command:
#      /usr/local/bin/train_tft_parquet \
#        --parquet-file /runpod-volume/test_data/ES_FUT_180d.parquet \
#        --epochs 50
#
# =============================================================================