# Multi-stage build for Foxhunt ML Training Service
FROM nvidia/cuda:12.1-devel-ubuntu22.04 as builder

# Install Rust and system dependencies
RUN apt-get update && apt-get install -y \
    curl \
    build-essential \
    pkg-config \
    libssl-dev \
    libpq-dev \
    protobuf-compiler \
    python3-dev \
    libblas-dev \
    liblapack-dev \
    && rm -rf /var/lib/apt/lists/*

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

# Set workspace directory
WORKDIR /workspace

# Copy workspace Cargo files
COPY ../Cargo.toml ../Cargo.lock ./
COPY ../trading_engine ./trading_engine
COPY ../risk ./risk
COPY ../data ./data
COPY ../ml ./ml

# Build the ML training service
RUN cargo build --release -p ml

# === RUNTIME IMAGE ===
FROM nvidia/cuda:12.1-runtime-ubuntu22.04

# Install runtime dependencies
RUN apt-get update && apt-get install -y \
    ca-certificates \
    libssl3 \
    libpq5 \
    curl \
    python3 \
    python3-pip \
    libblas3 \
    liblapack3 \
    && rm -rf /var/lib/apt/lists/*

# Install Python ML packages
RUN pip3 install \
    torch==2.1.0+cu121 \
    torchvision==0.16.0+cu121 \
    torchaudio==2.1.0+cu121 \
    --index-url https://download.pytorch.org/whl/cu121 \
    && pip3 install \
    tensorboard \
    numpy \
    pandas \
    scikit-learn \
    matplotlib \
    seaborn

# Create app user
RUN groupadd -r foxhunt && useradd -r -g foxhunt foxhunt

# Create directories
RUN mkdir -p /app/config /app/models /app/data /app/checkpoints /app/logs \
    && chown -R foxhunt:foxhunt /app

# Copy binary from builder
COPY --from=builder /workspace/target/release/ml /app/ml_service
RUN chmod +x /app/ml_service

# Copy configuration templates
COPY config/ /app/config/

USER foxhunt
WORKDIR /app

# Expose ports
EXPOSE 8082 6006 9002

# Health check
HEALTHCHECK --interval=30s --timeout=15s --start-period=120s --retries=3 \
    CMD curl -f http://localhost:8082/health || exit 1

# Set environment variables
ENV RUST_LOG=info
ENV FOXHUNT_CONFIG=/app/config/config.toml
ENV CUDA_VISIBLE_DEVICES=0
ENV NVIDIA_VISIBLE_DEVICES=0
ENV PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512

CMD ["./ml_service"]