FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 AS base
RUN apt-get update && \
    apt-get upgrade -y && \
    apt install -y \
    build-essential \
    curl \
    git \
    gnupg \
    jq \
    python3-dev \
    python3-venv \
    tmux \
    vim \
    wget

RUN wget -q -O - https://dl-ssl.google.com/linux/linux_signing_key.pub | gpg --dearmor -o /etc/apt/trusted.gpg.d/google.gpg
RUN echo "deb http://dl.google.com/linux/chrome/deb/ stable main" >> /etc/apt/sources.list.d/google-chrome.list
RUN apt-get update && \
    apt-get install -y \
    google-chrome-stable

RUN rm -rf /var/lib/apt/lists/*

RUN curl https://sh.rustup.rs -sSf | bash -s -- -y
ENV PATH="/root/.cargo/bin:${PATH}"

# Create directory structure
RUN mkdir -p \
    /root/sia \
    /root/sia/scripts \
    /root/data/iterations \
    /root/data/user \
    /root/data/tasks \
    /root/data/environment \
    /root/models \
    /root/desktop \
    /root/venvs

# ITB tool setup
FROM base AS itb-env
RUN python3 -m venv /root/venvs/itb --system-site-packages
COPY ./tools/itb/setup.py /root/sia/tools/itb/setup.py
RUN /root/venvs/itb/bin/python /root/sia/tools/itb/setup.py egg_info
RUN --mount=type=cache,target=/root/.cache/pip \
    /root/venvs/itb/bin/pip install -r *.egg-info/requires.txt
RUN rm -rf *.egg-info/

# Train tool setup
FROM base AS train-env
RUN python3 -m venv /root/venvs/train --system-site-packages
COPY ./tools/train/setup.py /root/sia/tools/train/setup.py
RUN /root/venvs/train/bin/python /root/sia/tools/train/setup.py egg_info
RUN --mount=type=cache,target=/root/.cache/pip \
    /root/venvs/train/bin/pip install -r *.egg-info/requires.txt
RUN rm -rf *.egg-info/

# SIA core setup
FROM base AS sia-env
RUN python3 -m venv /root/venvs/sia --system-site-packages
COPY ./setup.py /root/sia/setup.py
RUN /root/venvs/sia/bin/python /root/sia/setup.py egg_info
RUN sed -i '/\/root\/sia\/lib/d' *.egg-info/requires.txt
RUN --mount=type=cache,target=/root/.cache/pip \
    /root/venvs/sia/bin/pip install -r *.egg-info/requires.txt
RUN rm -rf *.egg-info/

# Notebook setup
FROM base AS notebook-env
RUN python3 -m venv /root/venvs/notebook --system-site-packages
RUN --mount=type=cache,target=/root/.cache/pip \
    /root/venvs/notebook/bin/pip install \
        jupyter \
        ipykernel \
        ipywidgets
RUN /root/venvs/notebook/bin/ipython kernel install --name=notebook

# Web frontend build
FROM base
# AS web-build
ENV NODE_VERSION=16.13.0
RUN curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.39.0/install.sh | bash
ENV NVM_DIR=/root/.nvm
RUN . "$NVM_DIR/nvm.sh" && nvm install ${NODE_VERSION}
RUN . "$NVM_DIR/nvm.sh" && nvm use v${NODE_VERSION}
RUN . "$NVM_DIR/nvm.sh" && nvm alias default v${NODE_VERSION}
ENV PATH="/root/.nvm/versions/node/v${NODE_VERSION}/bin/:${PATH}"

FROM node:20-alpine AS web-build
WORKDIR /app
COPY web/package.json ./
RUN npm install
RUN rm -rf /root/.npm/_cacache
COPY web .
RUN npm run build

# Final image
FROM base

COPY --from=itb-env /root/venvs/itb /root/venvs/itb
COPY --from=train-env /root/venvs/train /root/venvs/train
COPY --from=sia-env /root/venvs/sia /root/venvs/sia
COPY --from=notebook-env /root/venvs/notebook /root/venvs/notebook
COPY --from=web-build /app/dist /root/static/

RUN echo 'source /root/sia/scripts/add_venvs_to_path.sh' >> /root/.bashrc
RUN echo 'source /root/sia/scripts/add_venvs_to_path.sh' >> /etc/profile

WORKDIR /root/desktop

ENTRYPOINT ["/bin/bash", "-c"]
CMD ["/root/sia/scripts/install.sh; /bin/bash -lc /root/sia/scripts/restart.sh"]