ARG BASE_IMAGE=ubuntu:24.04
ARG CHALLENGE

# Development environment stage with CUDA support
FROM ${BASE_IMAGE}

ARG CHALLENGE
ENV CHALLENGE=${CHALLENGE}
RUN if [ -z "$CHALLENGE" ]; then echo "Error: '--build-arg CHALLENGE' must be set." && exit 1; fi

RUN ARCH=$(uname -m) && \
    if [ "$ARCH" != "aarch64" ] && [ "$ARCH" != "arm64" ] && [ "$ARCH" != "amd64" ] && [ "$ARCH" != "x86_64" ]; then \
        echo "Unsupported architecture: $ARCH. Must be 'aarch64', 'arm64', 'amd64', or 'x86_64'." && exit 1; \
    fi

RUN apt update && apt install -y curl build-essential zstd python3 python3-pip wget gnupg ca-certificates

# Install CUDA toolkit
RUN ARCH=$(uname -m) && \
    if [ "$ARCH" = "aarch64" ] || [ "$ARCH" = "arm64" ]; then \
        CUDA_ARCH="sbsa"; \
    else \
        CUDA_ARCH="x86_64"; \
    fi && \
    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/${CUDA_ARCH}/cuda-keyring_1.1-1_all.deb && \
    dpkg -i cuda-keyring_1.1-1_all.deb && \
    apt-get update && \
    DEBIAN_FRONTEND=noninteractive apt-get install -y \
        cuda-toolkit-12-6 \
        libcublas-12-6 \
        libcublas-dev-12-6 \
        libcudnn9-cuda-12 \
        libcudnn9-dev-cuda-12 && \
    rm cuda-keyring_1.1-1_all.deb && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

# Set CUDA environment variables
ENV CUDA_HOME=/usr/local/cuda-12.6
ENV CUDA_PATH=/usr/local/cuda-12.6
ENV CUDA_ROOT=/usr/local/cuda-12.6
ENV CUDA_TOOLKIT_ROOT_DIR=/usr/local/cuda-12.6
ENV PATH=${CUDA_HOME}/bin:${PATH}
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH}
ENV LIBRARY_PATH=${CUDA_HOME}/lib64:/usr/lib/x86_64-linux-gnu:${LIBRARY_PATH}

# Create symlinks for CUDA libraries
RUN ln -sf /usr/local/cuda-12.6 /usr/local/cuda && \
    ldconfig

RUN pip3 install blake3 requests --break-system-packages

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

RUN ARCH=$(uname -m) && \
    RUST_TARGET=$(if [ "$ARCH" = "aarch64" ] || [ "$ARCH" = "arm64" ]; then \
        echo "aarch64-unknown-linux-gnu"; \
    else \
        echo "x86_64-unknown-linux-gnu"; \
    fi) && \
    rustup install nightly-2025-02-10 && \
    rustup default nightly-2025-02-10 && \
    rustup component add rust-src && \
    rustup target add $RUST_TARGET && \
    RUST_LIBDIR=$(rustc --print target-libdir --target=$RUST_TARGET) && \
    ln -s $RUST_LIBDIR /usr/local/lib/rust
ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/lib/rust"

RUN ARCH=$(uname -m) && \
    LLVM_URL=$(if [ "$ARCH" = "aarch64" ] || [ "$ARCH" = "arm64" ]; then \
        echo "https://github.com/tig-foundation/llvm/releases/download/aarch64.rc4%2B19.1.7/llvm.tar.zst"; \
    else \
        echo "https://github.com/tig-foundation/llvm/releases/download/amd64.rc4%2B19.1.7/llvm.tar.zst"; \
    fi) && \
    curl -L $LLVM_URL -O && \
    mkdir -p /opt/llvm && \
    tar -xf llvm.tar.zst -C /opt/llvm && \
    rm -rf llvm.tar.zst && \
    ln -s /opt/llvm/bin/* /usr/local/bin/

COPY scripts/ /usr/local/bin/tig-scripts/
COPY tig-binary/scripts/ /usr/local/bin/tig-scripts/
RUN chmod +x /usr/local/bin/tig-scripts/*
ENV PATH="/usr/local/bin/tig-scripts:${PATH}"

COPY . /tmp/tig-monorepo
WORKDIR /tmp/tig-monorepo

RUN cargo build -r -p tig-runtime --features $CHALLENGE && \
    cargo build -r -p tig-verifier --features $CHALLENGE && \
    mv target/release/tig-runtime /usr/local/bin/ && \
    mv target/release/tig-verifier /usr/local/bin/ && \
    chmod +x /usr/local/bin/tig-runtime && \
    chmod +x /usr/local/bin/tig-verifier && \
    rm -rf tig-monorepo

WORKDIR /app