FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 AS base
RUN apt-get update && \
    apt-get upgrade -y && \
    apt install -y \
    python3-pip \
    git \
    jq \
    python3-venv \
    wget \
    gnupg \
    vim \
    curl
RUN wget -q -O - https://dl-ssl.google.com/linux/linux_signing_key.pub | apt-key add -
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/*

# 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
COPY ./scripts/setup_binaries.py /root/sia/scripts/
COPY ./tools/itb/setup.py /root/sia/tools/itb/setup.py
RUN python3 /root/sia/scripts/setup_binaries.py /root/sia/tools/itb/setup.py
RUN python3 -m venv /root/venvs/itb
RUN /root/venvs/itb/bin/pip install -e /root/sia/tools/itb/

# Train tool setup
FROM base AS train-env
COPY ./scripts/setup_binaries.py /root/sia/scripts/
COPY ./tools/train/setup.py /root/sia/tools/train/setup.py
RUN python3 /root/sia/scripts/setup_binaries.py /root/sia/tools/train/setup.py
RUN python3 -m venv /root/venvs/train
RUN /root/venvs/train/bin/pip install -e /root/sia/tools/train/

# SIA core setup
FROM base AS sia-env
COPY ./scripts/setup_binaries.py /root/sia/scripts/
COPY ./setup.py /root/sia/setup.py
RUN python3 /root/sia/scripts/setup_binaries.py /root/sia/setup.py
RUN python3 -m venv /root/venvs/sia
RUN /root/venvs/sia/bin/pip install -e /root/sia/

# Web frontend build
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=web-build /app/dist /root/static/

RUN for venv in /root/venvs/*/bin; do \
    echo "export PATH=\"$venv:\$PATH\"" >> /etc/profile.d/sia.sh ; \
done
RUN for venv in /root/venvs/*/bin; do \
    echo "export PATH=\"$venv:\$PATH\"" > "/etc/profile.d/venv_$(basename "$(dirname "$venv")").sh" \
done
WORKDIR /root/desktop

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