# Build stage
FROM nvidia/cuda:12.8.1-cudnn-runtime-ubuntu22.04 AS builder

# Install dependencies
RUN apt-get update && apt-get install -y \
    tar \
    curl \
    libatlas3-base \
    libblas3 \
    libjpeg8 \
    libpng16-16 \
    && rm -rf /var/lib/apt/lists/*

# Install Swift
RUN curl -O https://download.swift.org/swift-6.0.3-release/ubuntu2204/swift-6.0.3-RELEASE/swift-6.0.3-RELEASE-ubuntu22.04.tar.gz \
    && tar xzf swift-6.0.3-RELEASE-ubuntu22.04.tar.gz \
    && rm swift-6.0.3-RELEASE-ubuntu22.04.tar.gz \
    && mv swift-6.0.3-RELEASE-ubuntu22.04 /usr/local/swift

# Add Swift to PATH permanently
ENV PATH="/usr/local/swift/usr/bin:${PATH}"
ENV SWIFT_PATH="/usr/share/swift/usr/bin"
ENV LD_LIBRARY_PATH="/usr/lib/llvm-14/lib/:/usr/local/swift/usr/lib:/usr/local/swift/usr/lib/swift/linux/:${LD_LIBRARY_PATH}"
